distribution: add migration wire contracts

adapter/distribution_server.go

@@ -222,7 +222,7 @@ func (s *DistributionServer) saveSplitResultViaCoordinator(
 	}
 	nextRouteID := right.RouteID + 1
 
-	ops, err := buildCatalogSplitOps(parentID, left, right, nextVersion, nextRouteID)
+	ops, err := buildCatalogSplitOps(parentID, left, right, nextVersion, nextRouteID, s.catalog.AllowsRouteDescriptorV2Writes())
 	if err != nil {
 		return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "build split mutations: %v", err)
 	}
@@ -242,6 +242,7 @@ func buildCatalogSplitOps(
 	right distribution.RouteDescriptor,
 	nextVersion uint64,
 	nextRouteID uint64,
+	allowRouteDescriptorV2Writes bool,
 ) ([]*kv.Elem[kv.OP], error) {
 	// SplitRange mutates the catalog surgically: delete one parent route, add two
 	// children, bump the version, and advance the next-route-id counter.
@@ -251,7 +252,7 @@ func buildCatalogSplitOps(
 		Key: distribution.CatalogRouteKey(parentID),
 	})
 	for _, route := range []distribution.RouteDescriptor{left, right} {
-		encoded, err := distribution.EncodeRouteDescriptor(route)
+		encoded, err := distribution.EncodeRouteDescriptorForCatalogWrite(route, allowRouteDescriptorV2Writes)
 		if err != nil {
 			return nil, errors.WithStack(err)
 		}
@@ -384,20 +385,26 @@ func splitCatalogRoutes(
 ) (distribution.RouteDescriptor, distribution.RouteDescriptor) {
 	// parent and splitKey are already cloned before this point and are immutable here.
 	left := distribution.RouteDescriptor{
-		RouteID:       leftID,
-		Start:         parent.Start,
-		End:           splitKey,
-		GroupID:       parent.GroupID,
-		State:         parent.State,
-		ParentRouteID: parent.RouteID,
+		RouteID:                leftID,
+		Start:                  parent.Start,
+		End:                    splitKey,
+		GroupID:                parent.GroupID,
+		State:                  parent.State,
+		ParentRouteID:          parent.RouteID,
+		StagedVisibilityActive: parent.StagedVisibilityActive,
+		MigrationJobID:         parent.MigrationJobID,
+		MinWriteTSExclusive:    parent.MinWriteTSExclusive,
 	}
 	right := distribution.RouteDescriptor{
-		RouteID:       rightID,
-		Start:         splitKey,
-		End:           parent.End,
-		GroupID:       parent.GroupID,
-		State:         parent.State,
-		ParentRouteID: parent.RouteID,
+		RouteID:                rightID,
+		Start:                  splitKey,
+		End:                    parent.End,
+		GroupID:                parent.GroupID,
+		State:                  parent.State,
+		ParentRouteID:          parent.RouteID,
+		StagedVisibilityActive: parent.StagedVisibilityActive,
+		MigrationJobID:         parent.MigrationJobID,
+		MinWriteTSExclusive:    parent.MinWriteTSExclusive,
 	}
 	return left, right
 }
@@ -447,12 +454,15 @@ func toProtoRouteDescriptors(routes []distribution.RouteDescriptor) []*pb.RouteD
 
 func toProtoRouteDescriptor(route distribution.RouteDescriptor) *pb.RouteDescriptor {
 	return &pb.RouteDescriptor{
-		RouteId:       route.RouteID,
-		Start:         distribution.CloneBytes(route.Start),
-		End:           distribution.CloneBytes(route.End),
-		RaftGroupId:   route.GroupID,
-		State:         toProtoRouteState(route.State),
-		ParentRouteId: route.ParentRouteID,
+		RouteId:                route.RouteID,
+		Start:                  distribution.CloneBytes(route.Start),
+		End:                    distribution.CloneBytes(route.End),
+		RaftGroupId:            route.GroupID,
+		State:                  toProtoRouteState(route.State),
+		ParentRouteId:          route.ParentRouteID,
+		StagedVisibilityActive: route.StagedVisibilityActive,
+		MigrationJobId:         route.MigrationJobID,
+		MinWriteTsExclusive:    route.MinWriteTSExclusive,
 	}
 }
 

adapter/distribution_server_test.go

@@ -76,15 +76,18 @@ func TestDistributionServerListRoutes_ReadsDurableCatalog(t *testing.T) {
 	t.Parallel()
 
 	ctx := context.Background()
-	catalog := distribution.NewCatalogStore(store.NewMVCCStore())
+	catalog := distribution.NewCatalogStore(store.NewMVCCStore(), distribution.WithCatalogRouteDescriptorV2Writes(true))
 	saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{
 		{
-			RouteID:       2,
-			Start:         []byte("m"),
-			End:           nil,
-			GroupID:       2,
-			State:         distribution.RouteStateWriteFenced,
-			ParentRouteID: 1,
+			RouteID:                2,
+			Start:                  []byte("m"),
+			End:                    nil,
+			GroupID:                2,
+			State:                  distribution.RouteStateWriteFenced,
+			ParentRouteID:          1,
+			StagedVisibilityActive: true,
+			MigrationJobID:         42,
+			MinWriteTSExclusive:    99,
 		},
 		{
 			RouteID:       1,
@@ -111,6 +114,9 @@ func TestDistributionServerListRoutes_ReadsDurableCatalog(t *testing.T) {
 	require.Equal(t, uint64(2), resp.Routes[1].RouteId)
 	require.Nil(t, resp.Routes[1].End)
 	require.Equal(t, pb.RouteState_ROUTE_STATE_WRITE_FENCED, resp.Routes[1].State)
+	require.True(t, resp.Routes[1].StagedVisibilityActive)
+	require.Equal(t, uint64(42), resp.Routes[1].MigrationJobId)
+	require.Equal(t, uint64(99), resp.Routes[1].MinWriteTsExclusive)
 }
 
 func TestDistributionServerListRoutes_RequiresCatalog(t *testing.T) {
@@ -128,15 +134,16 @@ func TestDistributionServerSplitRange_Success(t *testing.T) {
 
 	ctx := context.Background()
 	baseStore := store.NewMVCCStore()
-	catalog := distribution.NewCatalogStore(baseStore)
+	catalog := distribution.NewCatalogStore(baseStore, distribution.WithCatalogRouteDescriptorV2Writes(true))
 	saved, err := catalog.Save(ctx, 0, []distribution.RouteDescriptor{
 		{
-			RouteID:       1,
-			Start:         []byte(""),
-			End:           []byte("m"),
-			GroupID:       1,
-			State:         distribution.RouteStateActive,
-			ParentRouteID: 0,
+			RouteID:             1,
+			Start:               []byte(""),
+			End:                 []byte("m"),
+			GroupID:             1,
+			State:               distribution.RouteStateActive,
+			ParentRouteID:       0,
+			MinWriteTSExclusive: 99,
 		},
 		{
 			RouteID:       2,
@@ -169,28 +176,34 @@ func TestDistributionServerSplitRange_Success(t *testing.T) {
 	require.Equal(t, []byte("g"), resp.Left.End)
 	require.Equal(t, uint64(1), resp.Left.RaftGroupId)
 	require.Equal(t, uint64(1), resp.Left.ParentRouteId)
+	require.Equal(t, uint64(99), resp.Left.MinWriteTsExclusive)
 	require.Equal(t, uint64(4), resp.Right.RouteId)
 	require.Equal(t, []byte("g"), resp.Right.Start)
 	require.Equal(t, []byte("m"), resp.Right.End)
 	require.Equal(t, uint64(1), resp.Right.RaftGroupId)
 	require.Equal(t, uint64(1), resp.Right.ParentRouteId)
+	require.Equal(t, uint64(99), resp.Right.MinWriteTsExclusive)
 
 	snapshot, err := catalog.Snapshot(ctx)
 	require.NoError(t, err)
 	require.Equal(t, uint64(2), snapshot.Version)
 	require.Len(t, snapshot.Routes, 3)
 	// Catalog snapshots are sorted by range start key.
 	require.Equal(t, uint64(3), snapshot.Routes[0].RouteID)
+	require.Equal(t, uint64(99), snapshot.Routes[0].MinWriteTSExclusive)
 	require.Equal(t, uint64(4), snapshot.Routes[1].RouteID)
+	require.Equal(t, uint64(99), snapshot.Routes[1].MinWriteTSExclusive)
 	require.Equal(t, uint64(2), snapshot.Routes[2].RouteID)
 
 	require.Equal(t, uint64(2), engine.Version())
 	leftRoute, ok := engine.GetRoute([]byte("b"))
 	require.True(t, ok)
 	require.Equal(t, uint64(3), leftRoute.RouteID)
+	require.Equal(t, uint64(99), leftRoute.MinWriteTSExclusive)
 	rightRoute, ok := engine.GetRoute([]byte("h"))
 	require.True(t, ok)
 	require.Equal(t, uint64(4), rightRoute.RouteID)
+	require.Equal(t, uint64(99), rightRoute.MinWriteTSExclusive)
 }
 
 func TestDistributionServerSplitRange_RequiresCoordinator(t *testing.T) {
@@ -524,7 +537,7 @@ func TestBuildCatalogSplitOps_UsesSurgicalSplitMutations(t *testing.T) {
 		ParentRouteID: 1,
 	}
 
-	ops, err := buildCatalogSplitOps(1, left, right, 2, 5)
+	ops, err := buildCatalogSplitOps(1, left, right, 2, 5, false)
 	require.NoError(t, err)
 	require.Len(t, ops, 5)
 	require.Equal(t, kv.Del, ops[0].Op)

adapter/grpc.go

@@ -34,6 +34,22 @@ type GRPCServer struct {
 	pb.UnimplementedTransactionalKVServer
 }
 
+type rawReadFenceGetter interface {
+	GetAtWithReadFence(ctx context.Context, key []byte, ts uint64, groupID uint64, readRouteVersion uint64) ([]byte, error)
+}
+
+type rawReadFenceCommitTSReader interface {
+	LatestCommitTSWithReadFence(ctx context.Context, key []byte, readRouteVersion uint64) (uint64, bool, error)
+}
+
+type rawReadFenceScanner interface {
+	ScanAtWithReadFence(ctx context.Context, start []byte, end []byte, limit int, ts uint64, reverse bool, groupID uint64, readRouteVersion uint64, routeStart []byte, routeEnd []byte) ([]*store.KVPair, error)
+}
+
+type rawReadFenceVersioner interface {
+	ReadRouteVersion() uint64
+}
+
 type GRPCServerOption func(*GRPCServer)
 
 type rawGroupGetter interface {
@@ -98,7 +114,9 @@ func (r *GRPCServer) RawGet(ctx context.Context, req *pb.RawGetRequest) (*pb.Raw
 
 	var v []byte
 	var err error
-	if groupID := req.GetGroupId(); groupID != 0 {
+	if fenceGetter, ok := r.store.(rawReadFenceGetter); ok {
+		v, err = fenceGetter.GetAtWithReadFence(ctx, req.Key, readTS, req.GetGroupId(), r.readRouteVersion(req.GetReadRouteVersion()))
+	} else if groupID := req.GetGroupId(); groupID != 0 {
 		groupGetter, ok := r.store.(rawGroupGetter)
 		if !ok {
 			return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "raw get with explicit group requires a group-aware store"))
@@ -137,7 +155,14 @@ func (r *GRPCServer) RawLatestCommitTS(ctx context.Context, req *pb.RawLatestCom
 		}, nil
 	}
 
-	ts, exists, err := r.store.LatestCommitTS(ctx, key)
+	var ts uint64
+	var exists bool
+	var err error
+	if fenceReader, ok := r.store.(rawReadFenceCommitTSReader); ok {
+		ts, exists, err = fenceReader.LatestCommitTSWithReadFence(ctx, key, r.readRouteVersion(req.GetReadRouteVersion()))
+	} else {
+		ts, exists, err = r.store.LatestCommitTS(ctx, key)
+	}
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
@@ -159,21 +184,7 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (*
 		readTS = globalSnapshotTS(ctx, r.clock(), r.store)
 	}
 
-	var res []*store.KVPair
-	if groupID := req.GetGroupId(); groupID != 0 {
-		if req.GetReverse() {
-			return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans"))
-		}
-		groupScanner, ok := r.store.(rawGroupScanner)
-		if !ok {
-			return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, "raw scan with explicit group requires a group-aware store"))
-		}
-		res, err = groupScanner.ScanGroupAt(ctx, groupID, req.StartKey, req.EndKey, limit, readTS)
-	} else if req.GetReverse() {
-		res, err = r.store.ReverseScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
-	} else {
-		res, err = r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
-	}
+	res, err := r.rawScanAt(ctx, req, limit, readTS)
 	if err != nil {
 		if errors.Is(err, store.ErrReadTSCompacted) {
 			return &pb.RawScanAtResponse{Kv: nil}, errors.WithStack(status.Error(codes.FailedPrecondition, store.ErrReadTSCompacted.Error()))
@@ -184,6 +195,62 @@ func (r *GRPCServer) RawScanAt(ctx context.Context, req *pb.RawScanAtRequest) (*
 	return &pb.RawScanAtResponse{Kv: rawKvPairs(res)}, nil
 }
 
+func (r *GRPCServer) rawScanAt(ctx context.Context, req *pb.RawScanAtRequest, limit int, readTS uint64) ([]*store.KVPair, error) {
+	if fenceScanner, ok := r.store.(rawReadFenceScanner); ok {
+		if req.GetGroupId() != 0 && req.GetReverse() && !req.GetRouteBoundsPresent() {
+			return nil, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans"))
+		}
+		routeStart, routeEnd := rawScanRouteBounds(req)
+		res, err := fenceScanner.ScanAtWithReadFence(ctx, req.StartKey, req.EndKey, limit, readTS, req.GetReverse(), req.GetGroupId(), r.readRouteVersion(req.GetReadRouteVersion()), routeStart, routeEnd)
+		return res, errors.WithStack(err)
+	}
+	if req.GetGroupId() != 0 && req.GetReverse() {
+		return nil, errors.WithStack(status.Error(codes.InvalidArgument, "raw scan with explicit group does not support reverse scans"))
+	}
+	if groupID := req.GetGroupId(); groupID != 0 {
+		groupScanner, ok := r.store.(rawGroupScanner)
+		if !ok {
+			return nil, errors.WithStack(status.Error(codes.FailedPrecondition, "raw scan with explicit group requires a group-aware store"))
+		}
+		res, err := groupScanner.ScanGroupAt(ctx, groupID, req.StartKey, req.EndKey, limit, readTS)
+		return res, errors.WithStack(err)
+	}
+	if req.GetReverse() {
+		res, err := r.store.ReverseScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
+		return res, errors.WithStack(err)
+	}
+	res, err := r.store.ScanAt(ctx, req.StartKey, req.EndKey, limit, readTS)
+	return res, errors.WithStack(err)
+}
+
+func rawScanRouteBounds(req *pb.RawScanAtRequest) ([]byte, []byte) {
+	if req == nil {
+		return nil, nil
+	}
+	routeStart := req.GetRouteStart()
+	routeEnd := req.GetRouteEnd()
+	if req.GetRouteBoundsPresent() {
+		if routeStart == nil {
+			routeStart = []byte{}
+		}
+		if routeEnd == nil {
+			routeEnd = []byte{}
+		}
+	}
+	return routeStart, routeEnd
+}
+
+func (r *GRPCServer) readRouteVersion(requested uint64) uint64 {
+	if requested != 0 {
+		return requested
+	}
+	versioner, ok := r.store.(rawReadFenceVersioner)
+	if !ok {
+		return 0
+	}
+	return versioner.ReadRouteVersion()
+}
+
 func rawScanLimit(limit64 int64) (int, error) {
 	if limit64 < 0 {
 		return 0, errors.WithStack(kv.ErrInvalidRequest)