Region allocation should be a CTE

[smklein]

See: #1683 , https://github.com/oxidecomputer/remote-access-preview/issues/23 for context.

Here's the txn under consideration:

omicron/nexus/src/db/datastore/region.rs

Line 153 in 9f9ab18

.transaction(move |conn| {

Filing this issue as a PSA that I'm working on it.

[smklein]

Okay, here's my proposal. @jmpesp , @davepacheco , I'd appreciate feedback on how this looks to y'all.
If it looks good, I can begin the process of converting it to a Rust-wrapped sequence of queries.

I suspect I'll need to expand this a bit to make the error cases more informative, but I believe the skeleton of this CTE should work.

WITH
  /*
   * For idempotency, identify if we've already completed the region allocation.
   */
  previously_allocated_regions AS (
    SELECT
      id, time_created, time_modified, dataset_id, volume_id, block_size, blocks_per_extent, extent_count
    FROM
      omicron.public.Region
    WHERE
      volume_id = <volume ID to be inserted>
  ),

  /*
   * The crux of our allocation function.
   * Find all the Crucible datasets we'd like to use
   * as targets on which to allocate the new regions.
   */
  candidate_datasets AS (
    SELECT
      omicron.public.Dataset.id,
      omicron.public.Dataset.time_created,
      omicron.public.Dataset.time_modified,
      omicron.public.Dataset.time_deleted,
      omicron.public.Dataset.rcgen,
      omicron.public.Dataset.pool_id,
      omicron.public.Dataset.ip,
      omicron.public.Dataset.port,
      omicron.public.Dataset.kind,
      omicron.public.Dataset.size_used
    FROM
      omicron.public.Dataset
    WHERE
      size_used IS NOT NULL AND
      time_deleted IS NULL AND
      kind = 'crucible'
    ORDER BY size_used ASC
    LIMIT <Region Redundancy Constant>
  ),

  /*
   * The zpools which contain the target datasets.
   * Used later to ensure we have enough space.
   */
  candidate_zpools AS (
    SELECT
      omicron.public.Zpool.id,
      omicron.public.Zpool.time_created,
      omicron.public.Zpool.time_modified,
      omicron.public.Zpool.time_deleted,
      omicron.public.Zpool.rcgen,
      omicron.public.Zpool.sled_id,
      omicron.public.Zpool.total_size
    FROM
      omicron.public.Zpool
    JOIN
      candidate_datasets
    ON
      candidate_datasets.pool_id = omicron.public.Zpool.id
  ),

  /*
   * Construct the regions we are trying to insert.
   */
  candidate_regions AS (
    SELECT
      gen_random_uuid() as id,
      now() as time_created,
      now() as time_modified,
      candidate_datasets.id as dataset_id,
      CAST(<volume ID to be inserted> AS UUID) as volume_id,
      <Block Size to be inserted> as block_size,
      <Blocks Per Extent to be inserted> as blocks_per_extent,
      <Extent Count to be inserted> as extent_count
    FROM
      candidate_datasets
  ),

  /*
   * Identify the delta in size we're proposing to the datasets (and zpools).
   */
  proposed_dataset_size_changes AS (
    SELECT
      pool_id,
      dataset_id,
      total_size as pool_total_size,
      block_size * blocks_per_extent * extent_count as size_used_delta
    FROM
      candidate_regions
    LEFT JOIN
      omicron.public.Dataset
    ON
      omicron.public.Dataset.id = candidate_regions.dataset_id
    LEFT JOIN
      omicron.public.Zpool
    ON
      omicron.public.Zpool.id = omicron.public.Dataset.pool_id
  ),

  /*
   * Get the total size used on all datasets in the candidate zpools.
   * This is a necessary calculation to determine if the proposed regions
   * will fit or not.
   */
  previous_zpool_size_usage AS (
    SELECT
      pool_id,
      SUM(size_used) as previous_size_used
    FROM
      omicron.public.Dataset
    LEFT JOIN
      candidate_zpools
    ON
      dataset.pool_id = candidate_zpools.id
    WHERE
      dataset.size_used IS NOT NULL AND
      dataset.time_deleted IS NULL
    GROUP BY
      pool_id
  ),

  /*
   * Compare the old space usage with the proposed changes,
   * and make a call on whether or not it'll fit.
   */
  proposed_zpool_size_usage AS (
    SELECT
      previous_size_used + size_used_delta <= pool_total_size as valid
    FROM
      previous_zpool_size_usage
    JOIN
      proposed_dataset_size_changes
    ON
      proposed_dataset_size_changes.pool_id = previous_zpool_size_usage.pool_id
  ),

  /*
   * Make a single decision: Should we do the insert/update, or not?
   */
  do_insert AS (
    SELECT IF(
      /* Only insert if we have not previously allocated regions. */
      NOT(EXISTS(SELECT id from previously_allocated_regions)) AND
      /* Only insert if we found the necessary amount of regions. */
      (SELECT COUNT(*) FROM candidate_regions) >= <Region Redundancy Constant> AND
      /* Only insert if all the proposed size changes fit. */
      (
        SELECT(BOOL_AND(valid))
        FROM proposed_zpool_size_usage
      ),
      TRUE,
      FALSE
    )
  ),

  /*
   * (if do_insert is true) Add the regions.
   */
  inserted_regions AS (
    INSERT INTO omicron.public.Region
      (
        SELECT * FROM candidate_regions
        WHERE (SELECT * FROM do_insert)
      )
    RETURNING *
  ),

  /*
   * (if do_insert is true) Update the datasets.
   */
  updated_datasets AS (
    UPDATE omicron.public.Dataset SET
      size_used =
        (
          size_used +
            (
              SELECT size_used_delta FROM proposed_dataset_size_changes
              WHERE proposed_dataset_size_changes.dataset_id = id
            )
        )
    WHERE
      (id IN (SELECT dataset_id FROM proposed_dataset_size_changes)) AND
      (SELECT * FROM do_insert)
    RETURNING *
  )
SELECT * FROM
  /* Return all old regions, if any existed */
  (SELECT * FROM previously_allocated_regions)
  /* Then return all new regions, if any were allocated */
  FULL OUTER JOIN (SELECT * FROM inserted_regions) ON TRUE;

[smklein]

For testing purposes: I've been using a local sql file called cte.sql to run experiments. It contains a variant of the above CTE.

• I run cockroach start-single-node --insecure in one terminal window
• I run cockroach sql -f cte.sql --insecure in another window

[smklein]

Some condensed feedback:

• If we can structure the INSERT clause to have an ON CONFLICT and appropriate RETURNING clause, we could avoid the up-front query for existing regions. This may or may not be doable, depending on the schema.
• It would be nice if the size checks were CONSTRAINTs on the DB. In particular, why are we calculating the diff of size on the whole zpool? If each dataset had a quota recorded on the table, we could compare the "size used" vs the "quota size". More detail in RFD 312, but it seems like this would let us avoid referencing the zpools altogether.
• Disk allocation should use multiple region sets #1644 was mentioned. Since this changes the allocation algorithm (requires using more+smaller regions than the "REDUNDANCY" constant) it would also change the CTE.