Infrastructure and Database Systems Questions

-- migration_control table (one row per migration)
CREATE TABLE IF NOT EXISTS migration_control (
  name text PRIMARY KEY,
  last_id bigint,
  status text,
  updated_at timestamptz DEFAULT now()
);

-- Backfill function: computes new_col using some expression compute_new_col(payload)
CREATE OR REPLACE FUNCTION backfill_new_col(batch_size int, migrate_name text)
RETURNS void AS $$
DECLARE
  got_lock boolean;
  ids bigint[];
  cur_id bigint;
BEGIN
  -- ensure control row exists
  INSERT INTO migration_control(name, last_id, status) VALUES (migrate_name, 0, 'running')
    ON CONFLICT (name) DO NOTHING;

  -- Acquire session advisory lock to prevent concurrent runners
  SELECT pg_try_advisory_lock(hashtext(migrate_name)) INTO got_lock;
  IF NOT got_lock THEN
    RAISE NOTICE 'Another backfill process is running for %', migrate_name;
    RETURN;
  END IF;

  LOOP
    -- pick a batch of ids that still need computation, > last_id to make progress monotonic
    SELECT last_id INTO cur_id FROM migration_control WHERE name = migrate_name FOR UPDATE;

    SELECT array_agg(id) INTO ids
    FROM (
      SELECT id
      FROM my_table
      WHERE id > COALESCE(cur_id,0) AND new_col IS NULL
      ORDER BY id
      FOR UPDATE SKIP LOCKED
      LIMIT batch_size
    ) t;

    IF ids IS NULL OR array_length(ids,1) = 0 THEN
      -- mark complete and exit
      UPDATE migration_control SET status='completed', updated_at=now() WHERE name = migrate_name;
      EXIT;
    END IF;

    BEGIN
      -- update batch in a short transaction
      UPDATE my_table m
      SET new_col = compute_new_col(m.payload)  -- replace with actual expression/function
      WHERE m.id = ANY(ids);

      -- persist progress: set last_id to the max id processed
      UPDATE migration_control
      SET last_id = (SELECT max(id) FROM unnest(ids) AS id), updated_at=now()
      WHERE name = migrate_name;

      -- commit implicitly by end of block
    EXCEPTION WHEN others THEN
      -- on error, log and retry later; don't advance last_id so resume picks up unprocessed rows
      RAISE WARNING 'Batch update failed: %', SQLERRM;
      PERFORM pg_sleep(1); -- simple backoff; in real code use exponential backoff
    END;

    -- small pause to reduce IO contention
    PERFORM pg_sleep(0.05);
  END LOOP;

  PERFORM pg_advisory_unlock(hashtext(migrate_name));
END;
$$ LANGUAGE plpgsql;

import hashlib
import time
from enum import Enum

class Phase(Enum): LOCKED=1; COPYING=2; SWITCHED=3; DRAINED=4; FINAL=5

# Durable metadata store (e.g., ZK/etcd)
meta = DurableKV()

def checksum(data):
    return hashlib.sha256(data).hexdigest()

def migrate_range(range_id, start, end, src, dst, chunk_size=1024):
    # idempotent read of phase
    phase = meta.get(f"{range_id}/phase") or None

    # 1) Acquire range lock (lease)
    if phase is None:
        if not meta.lease_acquire(f"{range_id}/lock", holder="coord"):
            raise Exception("lock busy")
        meta.put(f"{range_id}/phase", Phase.LOCKED.name)

    # 2) Copy data in chunks with checksums, resumable
    meta.put_if_not_exists(f"{range_id}/copy_pos", start)
    pos = meta.get(f"{range_id}/copy_pos")
    while pos < end:
        chunk = src.read_range(pos, min(pos+chunk_size, end))
        cs = checksum(chunk)
        dst.write_at(pos, chunk)                 # write to target
        dst.write_metadata(pos, cs)              # store checksum per chunk
        # verify read-back (optional)
        if dst.read_checksum(pos, len(chunk)) != cs:
            raise Exception("checksum mismatch")
        pos += len(chunk)
        meta.put(f"{range_id}/copy_pos", pos)    # durable progress marker
    meta.put(f"{range_id}/phase", Phase.COPYING.name)

    # 3) Atomic routing switch: use consensus update
    # write planned routing entry; other routers read routing atomically from meta
    meta.txn_put({f"{range_id}/route": dst.id, f"{range_id}/route_version": time.time()})
    meta.put(f"{range_id}/phase", Phase.SWITCHED.name)

    # 4) Drain in-flight/writes: enable forwarding on source for writes after switch
    # mark source to forward writes; wait until outstanding inflight count == 0
    src.enable_forwarding(range_id, dst.address)
    timeout = time.time() + 30
    while src.inflight_writes(range_id) > 0 and time.time() < timeout:
        time.sleep(0.1)
    if src.inflight_writes(range_id) > 0:
        raise Exception("drain timeout")
    meta.put(f"{range_id}/phase", Phase.DRAINED.name)

    # 5) Finalize: remove keys from source (garbage collect) after safe window
    src.delete_range(start, end)
    meta.put(f"{range_id}/phase", Phase.FINAL.name)
    meta.lease_release(f"{range_id}/lock")
    return True