Move key ranges via postgres_fdw

coordinator/provider/coordinator.go

@@ -3,6 +3,7 @@ package provider
 import (
 	"context"
 	"crypto/tls"
+	"fmt"
 	"net"
 	"time"
 
@@ -323,9 +324,9 @@ func (qc *qdbCoordinator) lockCoordinator(ctx context.Context, initialRouter boo
 	return registerRouter()
 }
 
-// TODO : unit tests
 // RunCoordinator side effect: it runs an asynchronous goroutine
 // that checks the availability of the SPQR router
+// TODO : unit tests
 func (qc *qdbCoordinator) RunCoordinator(ctx context.Context, initialRouter bool) {
 	if !qc.lockCoordinator(ctx, initialRouter) {
 		return
@@ -335,32 +336,39 @@ func (qc *qdbCoordinator) RunCoordinator(ctx context.Context, initialRouter bool
 	if err != nil {
 		spqrlog.Zero.Error().
 			Err(err).
-			Msg("faild to list key ranges")
+			Msg("failed to list key ranges")
 	}
 
+	// Finish any key range move or data transfer transaction in progress
 	for _, r := range ranges {
+		move, err := qc.GetKeyRangeMove(context.TODO(), r.KeyRangeID)
+		if err != nil {
+			spqrlog.Zero.Error().Err(err).Msg("error getting key range move from qdb")
+		}
+
 		tx, err := qc.db.GetTransferTx(context.TODO(), r.KeyRangeID)
 		if err != nil {
-			continue
+			spqrlog.Zero.Error().Err(err).Msg("error getting data transfer transaction from qdb")
 		}
-		if tx.ToStatus == qdb.Commited && tx.FromStatus != qdb.Commited {
-			datatransfers.ResolvePreparedTransaction(context.TODO(), tx.FromShardId, tx.FromTxName, true)
-			tem := kr.MoveKeyRange{
-				ShardId: tx.ToShardId,
-				Krid:    r.KeyRangeID,
+
+		var krm *kr.MoveKeyRange
+		if move != nil {
+			krm = &kr.MoveKeyRange{
+				Krid:    move.KeyRangeID,
+				ShardId: move.ShardId,
 			}
-			err = qc.Move(context.TODO(), &tem)
-			if err != nil {
-				spqrlog.Zero.Error().Err(err).Msg("failed to move key range")
+		} else if tx != nil {
+			krm = &kr.MoveKeyRange{
+				Krid:    r.KeyRangeID,
+				ShardId: tx.ToShardId,
 			}
-		} else if tx.FromStatus != qdb.Commited {
-			datatransfers.ResolvePreparedTransaction(context.TODO(), tx.ToShardId, tx.ToTxName, false)
-			datatransfers.ResolvePreparedTransaction(context.TODO(), tx.FromShardId, tx.FromTxName, false)
 		}
 
-		err = qc.db.RemoveTransferTx(context.TODO(), r.KeyRangeID)
-		if err != nil {
-			spqrlog.Zero.Error().Err(err).Msg("error removing from qdb")
+		if krm != nil {
+			err = qc.Move(context.TODO(), krm)
+			if err != nil {
+				spqrlog.Zero.Error().Err(err).Msg("error moving key range")
+			}
 		}
 	}
 
@@ -796,8 +804,26 @@ func (qc *qdbCoordinator) RecordKeyRangeMove(ctx context.Context, m *qdb.MoveKey
 	return m.MoveId, nil
 }
 
+func (qc *qdbCoordinator) GetKeyRangeMove(ctx context.Context, krId string) (*qdb.MoveKeyRange, error) {
+	ls, err := qc.db.ListKeyRangeMoves(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, krm := range ls {
+		// after the coordinator restarts, it will continue the move that was previously initiated.
+		// key range move already exist for this key range
+		// complete it first
+		if krm.KeyRangeID == krId {
+			return krm, nil
+		}
+	}
+
+	return nil, nil
+}
+
 // Move key range from one logical shard to another
-// This function reshards data by locking a portion of it,
+// This function re-shards data by locking a portion of it,
 // making it unavailable for read and write access during the process.
 // TODO : unit tests
 func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error {
@@ -809,7 +835,7 @@ func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error
 		Str("shard-id", req.ShardId).
 		Msg("qdb coordinator move key range")
 
-	keyRange, err := qc.db.GetKeyRange(ctx, req.Krid)
+	keyRange, err := qc.GetKeyRange(ctx, req.Krid)
 	if err != nil {
 		return err
 	}
@@ -819,63 +845,91 @@ func (qc *qdbCoordinator) Move(ctx context.Context, req *kr.MoveKeyRange) error
 		return nil
 	}
 
-	moveId, err := qc.RecordKeyRangeMove(ctx,
-		&qdb.MoveKeyRange{
-			MoveId:     uuid.New().String(),
-			ShardId:    req.ShardId,
-			KeyRangeID: req.Krid,
-			Status:     qdb.MoveKeyRangePlanned,
-		})
-
-	/* NOOP if key range move was already recorded */
+	move, err := qc.GetKeyRangeMove(ctx, req.Krid)
 	if err != nil {
 		return err
 	}
 
-	krmv, err := qc.LockKeyRange(ctx, req.Krid)
-	if err != nil {
-		return err
-	}
-	defer func() {
-		if err := qc.UnlockKeyRange(ctx, req.Krid); err != nil {
-			spqrlog.Zero.Error().Err(err).Msg("")
+	if move == nil {
+		// No key range moves in progress
+		move = &qdb.MoveKeyRange{
+			MoveId:     uuid.New().String(),
+			ShardId:    req.ShardId,
+			KeyRangeID: req.Krid,
+			Status:     qdb.MoveKeyRangePlanned,
 		}
-	}()
-
-	defer func() {
-		// set compelted status in the end of key range move operation
-		if err := qc.db.UpdateKeyRangeMoveStatus(ctx, moveId, qdb.MoveKeyRangeComplete); err != nil {
-			spqrlog.Zero.Error().Err(err).Msg("")
+		_, err = qc.RecordKeyRangeMove(ctx, move)
+		if err != nil {
+			return err
 		}
-	}()
-
-	/* physical changes on shards */
-	err = datatransfers.MoveKeys(ctx, keyRange.ShardID, req.ShardId, *keyRange, qc.db)
-	if err != nil {
-		spqrlog.Zero.Error().Err(err).Msg("failed to move rows")
-		return err
 	}
 
-	// Update the state of the distributed key-range metadata
-	krmv.ShardID = req.ShardId
-	if err := ops.ModifyKeyRangeWithChecks(ctx, qc.db, krmv); err != nil {
-		// TODO: check if unlock here is ok
-		return err
-	}
+	for move != nil {
+		switch move.Status {
+		case qdb.MoveKeyRangePlanned:
+			// lock the key range
+			_, err = qc.LockKeyRange(ctx, req.Krid)
+			if err != nil {
+				return err
+			}
+			if err = qc.db.UpdateKeyRangeMoveStatus(ctx, move.MoveId, qdb.MoveKeyRangeComplete); err != nil {
+				return err
+			}
+			move.Status = qdb.MoveKeyRangeStarted
+		case qdb.MoveKeyRangeStarted:
+			// move the data
+			ds, err := qc.GetDistribution(ctx, keyRange.Distribution)
+			if err != nil {
+				return err
+			}
+			err = datatransfers.MoveKeys(ctx, keyRange.ShardID, req.ShardId, keyRange, ds, qc.db, qc)
+			if err != nil {
+				spqrlog.Zero.Error().Err(err).Msg("failed to move rows")
+				return err
+			}
 
-	// Notify all routers about scheme changes.
-	if err := qc.traverseRouters(ctx, func(cc *grpc.ClientConn) error {
-		cl := routerproto.NewKeyRangeServiceClient(cc)
-		moveResp, err := cl.MoveKeyRange(ctx, &routerproto.MoveKeyRangeRequest{
-			Id:        krmv.ID,
-			ToShardId: krmv.ShardID,
-		})
-		spqrlog.Zero.Debug().
-			Interface("response", moveResp).
-			Msg("move key range response")
-		return err
-	}); err != nil {
-		return err
+			// update key range
+			krg, err := qc.GetKeyRange(ctx, req.Krid)
+			if err != nil {
+				return err
+			}
+			krg.ShardID = req.ShardId
+			if err := ops.ModifyKeyRangeWithChecks(ctx, qc.db, krg); err != nil {
+				// TODO: check if unlock here is ok
+				return err
+			}
+
+			// Notify all routers about scheme changes.
+			if err := qc.traverseRouters(ctx, func(cc *grpc.ClientConn) error {
+				cl := routerproto.NewKeyRangeServiceClient(cc)
+				moveResp, err := cl.MoveKeyRange(ctx, &routerproto.MoveKeyRangeRequest{
+					Id:        krg.ID,
+					ToShardId: krg.ShardID,
+				})
+				spqrlog.Zero.Debug().
+					Interface("response", moveResp).
+					Msg("move key range response")
+				return err
+			}); err != nil {
+				return err
+			}
+
+			if err := qc.db.UpdateKeyRangeMoveStatus(ctx, move.MoveId, qdb.MoveKeyRangeComplete); err != nil {
+				spqrlog.Zero.Error().Err(err).Msg("")
+			}
+			move.Status = qdb.MoveKeyRangeComplete
+		case qdb.MoveKeyRangeComplete:
+			// unlock key range
+			if err := qc.UnlockKeyRange(ctx, req.Krid); err != nil {
+				spqrlog.Zero.Error().Err(err).Msg("")
+			}
+			if err := qc.db.DeleteKeyRangeMove(ctx, move.MoveId); err != nil {
+				return err
+			}
+			move = nil
+		default:
+			return fmt.Errorf("unknown key range move status: \"%s\"", move.Status)
+		}
 	}
 	return nil
 }

docker/shard/bin/setup

@@ -40,14 +40,15 @@ sudo -u postgres psql -p 6432 -h localhost -U postgres -c "CREATE ROLE $POSTGRES
 
 # Create databases
 sudo -u postgres psql -p 6432 -h localhost -U postgres -c "CREATE DATABASE $POSTGRES_DB" -d postgres >> $SETUP_LOG 2>&1 || {
-    echo "ERROR: users creation failed, examine the log"
+    echo "ERROR: databases creation failed, examine the log"
     cat "$SETUP_LOG"
     exit 1
 }
 
- # Create extension
- sudo -u postgres psql -p 6432 -h localhost -U postgres -c "CREATE EXTENSION IF NOT EXISTS pg_stat_statements; CREATE EXTENSION IF NOT EXISTS pg_stat_kcache; CREATE EXTENSION IF NOT EXISTS pg_comment_stats;" -d $POSTGRES_DB >> $SETUP_LOG 2>&1 || {
-     echo "ERROR: users creation failed, examine the log"
+
+# # Create extension
+ sudo -u postgres psql -p 6432 -h localhost -U postgres -c "CREATE EXTENSION IF NOT EXISTS postgres_fdw; CREATE EXTENSION IF NOT EXISTS pg_stat_statements; CREATE EXTENSION IF NOT EXISTS pg_stat_kcache; CREATE EXTENSION IF NOT EXISTS pg_comment_stats;" -d $POSTGRES_DB >> $SETUP_LOG 2>&1 || {
+     echo "ERROR: extensions creation failed, examine the log"
      cat "$SETUP_LOG"
      exit 1
  }