#include "scrub/scrub.h"
 #include "scrub/common.h"
 #include "scrub/trace.h"
+#include "scrub/btree.h"
 
 /* Common code for the metadata scrubbers. */
 
        trace_xfs_scrub_incomplete(sc, __return_address);
 }
 
+/*
+ * AG scrubbing
+ *
+ * These helpers facilitate locking an allocation group's header
+ * buffers, setting up cursors for all btrees that are present, and
+ * cleaning everything up once we're through.
+ */
+
+/*
+ * Grab all the headers for an AG.
+ *
+ * The headers should be released by xfs_scrub_ag_free, but as a fail
+ * safe we attach all the buffers we grab to the scrub transaction so
+ * they'll all be freed when we cancel it.
+ */
+int
+xfs_scrub_ag_read_headers(
+       struct xfs_scrub_context        *sc,
+       xfs_agnumber_t                  agno,
+       struct xfs_buf                  **agi,
+       struct xfs_buf                  **agf,
+       struct xfs_buf                  **agfl)
+{
+       struct xfs_mount                *mp = sc->mp;
+       int                             error;
+
+       error = xfs_ialloc_read_agi(mp, sc->tp, agno, agi);
+       if (error)
+               goto out;
+
+       error = xfs_alloc_read_agf(mp, sc->tp, agno, 0, agf);
+       if (error)
+               goto out;
+       if (!*agf) {
+               error = -ENOMEM;
+               goto out;
+       }
+
+       error = xfs_alloc_read_agfl(mp, sc->tp, agno, agfl);
+       if (error)
+               goto out;
+
+out:
+       return error;
+}
+
+/* Release all the AG btree cursors. */
+void
+xfs_scrub_ag_btcur_free(
+       struct xfs_scrub_ag             *sa)
+{
+       if (sa->refc_cur)
+               xfs_btree_del_cursor(sa->refc_cur, XFS_BTREE_ERROR);
+       if (sa->rmap_cur)
+               xfs_btree_del_cursor(sa->rmap_cur, XFS_BTREE_ERROR);
+       if (sa->fino_cur)
+               xfs_btree_del_cursor(sa->fino_cur, XFS_BTREE_ERROR);
+       if (sa->ino_cur)
+               xfs_btree_del_cursor(sa->ino_cur, XFS_BTREE_ERROR);
+       if (sa->cnt_cur)
+               xfs_btree_del_cursor(sa->cnt_cur, XFS_BTREE_ERROR);
+       if (sa->bno_cur)
+               xfs_btree_del_cursor(sa->bno_cur, XFS_BTREE_ERROR);
+
+       sa->refc_cur = NULL;
+       sa->rmap_cur = NULL;
+       sa->fino_cur = NULL;
+       sa->ino_cur = NULL;
+       sa->bno_cur = NULL;
+       sa->cnt_cur = NULL;
+}
+
+/* Initialize all the btree cursors for an AG. */
+int
+xfs_scrub_ag_btcur_init(
+       struct xfs_scrub_context        *sc,
+       struct xfs_scrub_ag             *sa)
+{
+       struct xfs_mount                *mp = sc->mp;
+       xfs_agnumber_t                  agno = sa->agno;
+
+       if (sa->agf_bp) {
+               /* Set up a bnobt cursor for cross-referencing. */
+               sa->bno_cur = xfs_allocbt_init_cursor(mp, sc->tp, sa->agf_bp,
+                               agno, XFS_BTNUM_BNO);
+               if (!sa->bno_cur)
+                       goto err;
+
+               /* Set up a cntbt cursor for cross-referencing. */
+               sa->cnt_cur = xfs_allocbt_init_cursor(mp, sc->tp, sa->agf_bp,
+                               agno, XFS_BTNUM_CNT);
+               if (!sa->cnt_cur)
+                       goto err;
+       }
+
+       /* Set up a inobt cursor for cross-referencing. */
+       if (sa->agi_bp) {
+               sa->ino_cur = xfs_inobt_init_cursor(mp, sc->tp, sa->agi_bp,
+                                       agno, XFS_BTNUM_INO);
+               if (!sa->ino_cur)
+                       goto err;
+       }
+
+       /* Set up a finobt cursor for cross-referencing. */
+       if (sa->agi_bp && xfs_sb_version_hasfinobt(&mp->m_sb)) {
+               sa->fino_cur = xfs_inobt_init_cursor(mp, sc->tp, sa->agi_bp,
+                               agno, XFS_BTNUM_FINO);
+               if (!sa->fino_cur)
+                       goto err;
+       }
+
+       /* Set up a rmapbt cursor for cross-referencing. */
+       if (sa->agf_bp && xfs_sb_version_hasrmapbt(&mp->m_sb)) {
+               sa->rmap_cur = xfs_rmapbt_init_cursor(mp, sc->tp, sa->agf_bp,
+                               agno);
+               if (!sa->rmap_cur)
+                       goto err;
+       }
+
+       /* Set up a refcountbt cursor for cross-referencing. */
+       if (sa->agf_bp && xfs_sb_version_hasreflink(&mp->m_sb)) {
+               sa->refc_cur = xfs_refcountbt_init_cursor(mp, sc->tp,
+                               sa->agf_bp, agno, NULL);
+               if (!sa->refc_cur)
+                       goto err;
+       }
+
+       return 0;
+err:
+       return -ENOMEM;
+}
+
+/* Release the AG header context and btree cursors. */
+void
+xfs_scrub_ag_free(
+       struct xfs_scrub_context        *sc,
+       struct xfs_scrub_ag             *sa)
+{
+       xfs_scrub_ag_btcur_free(sa);
+       if (sa->agfl_bp) {
+               xfs_trans_brelse(sc->tp, sa->agfl_bp);
+               sa->agfl_bp = NULL;
+       }
+       if (sa->agf_bp) {
+               xfs_trans_brelse(sc->tp, sa->agf_bp);
+               sa->agf_bp = NULL;
+       }
+       if (sa->agi_bp) {
+               xfs_trans_brelse(sc->tp, sa->agi_bp);
+               sa->agi_bp = NULL;
+       }
+       sa->agno = NULLAGNUMBER;
+}
+
+/*
+ * For scrub, grab the AGI and the AGF headers, in that order.  Locking
+ * order requires us to get the AGI before the AGF.  We use the
+ * transaction to avoid deadlocking on crosslinked metadata buffers;
+ * either the caller passes one in (bmap scrub) or we have to create a
+ * transaction ourselves.
+ */
+int
+xfs_scrub_ag_init(
+       struct xfs_scrub_context        *sc,
+       xfs_agnumber_t                  agno,
+       struct xfs_scrub_ag             *sa)
+{
+       int                             error;
+
+       sa->agno = agno;
+       error = xfs_scrub_ag_read_headers(sc, agno, &sa->agi_bp,
+                       &sa->agf_bp, &sa->agfl_bp);
+       if (error)
+               return error;
+
+       return xfs_scrub_ag_btcur_init(sc, sa);
+}
+
 /* Per-scrubber setup functions */
 
 /* Set us up with a transaction and an empty context. */