store: add migration version import export

adapter/distribution_server.go

@@ -10,6 +10,7 @@ import (
 	"github.com/bootjp/elastickv/distribution"
 	"github.com/bootjp/elastickv/kv"
 	pb "github.com/bootjp/elastickv/proto"
+	"github.com/bootjp/elastickv/store"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -22,6 +23,7 @@ type DistributionServer struct {
 	catalog     *distribution.CatalogStore
 	coordinator kv.Coordinator
 	readTracker *kv.ActiveTimestampTracker
+	readBlocked func() bool
 	reloadRetry struct {
 		attempts int
 		interval time.Duration
@@ -46,6 +48,12 @@ func WithDistributionActiveTimestampTracker(tracker *kv.ActiveTimestampTracker)
 	}
 }
 
+func WithDistributionReadGate(blocked func() bool) DistributionServerOption {
+	return func(s *DistributionServer) {
+		s.readBlocked = blocked
+	}
+}
+
 // WithCatalogReloadRetryPolicy configures the retry policy used after split
 // commit when waiting for the local catalog snapshot to become visible.
 func WithCatalogReloadRetryPolicy(attempts int, interval time.Duration) DistributionServerOption {
@@ -95,13 +103,30 @@ func NewDistributionServer(e *distribution.Engine, catalog *distribution.Catalog
 	return s
 }
 
+func (s *DistributionServer) SetReadGate(blocked func() bool) {
+	if s != nil {
+		s.readBlocked = blocked
+	}
+}
+
+func (s *DistributionServer) requireReadReady() error {
+	if s != nil && s.readBlocked != nil && s.readBlocked() {
+		//nolint:wrapcheck // Preserve the gRPC status code for startup readers.
+		return status.Error(codes.Unavailable, "distribution startup has not completed")
+	}
+	return nil
+}
+
 // UpdateRoute allows updating route information.
 func (s *DistributionServer) UpdateRoute(start, end []byte, group uint64) {
 	s.engine.UpdateRoute(start, end, group)
 }
 
 // GetRoute returns route for a key.
 func (s *DistributionServer) GetRoute(ctx context.Context, req *pb.GetRouteRequest) (*pb.GetRouteResponse, error) {
+	if err := s.requireReadReady(); err != nil {
+		return nil, err
+	}
 	r, ok := s.engine.GetRoute(req.Key)
 	if !ok {
 		return &pb.GetRouteResponse{}, nil
@@ -121,6 +146,9 @@ func (s *DistributionServer) GetTimestamp(ctx context.Context, req *pb.GetTimest
 
 // ListRoutes returns all durable routes from catalog storage.
 func (s *DistributionServer) ListRoutes(ctx context.Context, req *pb.ListRoutesRequest) (*pb.ListRoutesResponse, error) {
+	if err := s.requireReadReady(); err != nil {
+		return nil, err
+	}
 	snapshot, err := s.loadCatalogSnapshot(ctx)
 	if err != nil {
 		return nil, err
@@ -162,14 +190,18 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR
 	if err := validateSplitKey(parent, splitKey); err != nil {
 		return nil, err
 	}
+	splitJobReadKeys, err := s.splitJobOverlapReadKeys(ctx, snapshot, parent)
+	if err != nil {
+		return nil, err
+	}
 
 	leftID, rightID, err := s.allocateChildRouteIDs(ctx, snapshot.ReadTS, snapshot.Routes)
 	if err != nil {
 		return nil, err
 	}
 	left, right := splitCatalogRoutes(parent, splitKey, leftID, rightID)
 
-	saved, err := s.saveSplitResultViaCoordinator(ctx, snapshot.ReadTS, req.GetExpectedCatalogVersion(), parent.RouteID, left, right)
+	saved, err := s.saveSplitResultViaCoordinator(ctx, snapshot.ReadTS, req.GetExpectedCatalogVersion(), parent.RouteID, splitJobReadKeys, left, right)
 	if err != nil {
 		return nil, err
 	}
@@ -210,6 +242,7 @@ func (s *DistributionServer) saveSplitResultViaCoordinator(
 	readTS uint64,
 	expectedVersion uint64,
 	parentID uint64,
+	readKeys [][]byte,
 	left distribution.RouteDescriptor,
 	right distribution.RouteDescriptor,
 ) (distribution.CatalogSnapshot, error) {
@@ -227,10 +260,14 @@ func (s *DistributionServer) saveSplitResultViaCoordinator(
 		return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "build split mutations: %v", err)
 	}
 	if _, err := s.coordinator.Dispatch(ctx, &kv.OperationGroup[kv.OP]{
-		Elems:   ops,
-		IsTxn:   true,
-		StartTS: readTS,
+		Elems:    ops,
+		IsTxn:    true,
+		StartTS:  readTS,
+		ReadKeys: readKeys,
 	}); err != nil {
+		if errors.Is(err, store.ErrWriteConflict) {
+			return distribution.CatalogSnapshot{}, grpcStatusError(codes.Aborted, errDistributionCatalogConflict.Error())
+		}
 		return distribution.CatalogSnapshot{}, grpcStatusErrorf(codes.Internal, "commit split mutations: %v", err)
 	}
 	return s.loadCatalogSnapshotAtLeastVersion(ctx, nextVersion)
@@ -377,6 +414,81 @@ func validateSplitKey(parent distribution.RouteDescriptor, splitKey []byte) erro
 	return nil
 }
 
+func (s *DistributionServer) splitJobOverlapReadKeys(ctx context.Context, snapshot distribution.CatalogSnapshot, parent distribution.RouteDescriptor) ([][]byte, error) {
+	jobs, err := s.catalog.ListSplitJobsAt(ctx, snapshot.ReadTS)
+	if err != nil {
+		return nil, grpcStatusErrorf(codes.Internal, "load split jobs: %v", err)
+	}
+	readKeys := splitJobReadFenceKeys(jobs)
+	for _, job := range jobs {
+		if !splitJobIsLive(job) {
+			continue
+		}
+		for _, interval := range liveSplitJobIntervals(job, snapshot.Routes) {
+			if routeRangeIntersects(parent.Start, parent.End, interval.start, interval.end) {
+				return nil, grpcStatusError(codes.Aborted, distribution.ErrSplitJobOverlap.Error())
+			}
+		}
+	}
+	return readKeys, nil
+}
+
+func splitJobReadFenceKeys(jobs []distribution.SplitJob) [][]byte {
+	readKeys := make([][]byte, 0, len(jobs)+1)
+	readKeys = append(readKeys, distribution.CatalogNextSplitJobIDKey())
+	for _, job := range jobs {
+		if splitJobIsLive(job) {
+			readKeys = append(readKeys, distribution.CatalogSplitJobKey(job.JobID))
+		}
+	}
+	return readKeys
+}
+
+func splitJobIsLive(job distribution.SplitJob) bool {
+	return job.Phase != distribution.SplitJobPhaseDone && job.Phase != distribution.SplitJobPhaseAbandoned
+}
+
+type routeInterval struct {
+	start []byte
+	end   []byte
+}
+
+const initialLiveSplitJobIntervalCapacity = 2
+
+func liveSplitJobIntervals(job distribution.SplitJob, routes []distribution.RouteDescriptor) []routeInterval {
+	out := make([]routeInterval, 0, initialLiveSplitJobIntervalCapacity)
+	for _, route := range routes {
+		switch {
+		case route.RouteID == job.SourceRouteID:
+			out = append(out, routeInterval{
+				start: distribution.CloneBytes(job.SplitKey),
+				end:   distribution.CloneBytes(route.End),
+			})
+		case route.ParentRouteID == job.SourceRouteID && routeRangeIntersects(route.Start, route.End, job.SplitKey, nil):
+			out = append(out, routeInterval{
+				start: distribution.CloneBytes(route.Start),
+				end:   distribution.CloneBytes(route.End),
+			})
+		case job.JobID != 0 && route.MigrationJobID == job.JobID:
+			out = append(out, routeInterval{
+				start: distribution.CloneBytes(route.Start),
+				end:   distribution.CloneBytes(route.End),
+			})
+		}
+	}
+	return out
+}
+
+func routeRangeIntersects(aStart, aEnd, bStart, bEnd []byte) bool {
+	if aEnd != nil && bytes.Compare(aEnd, bStart) <= 0 {
+		return false
+	}
+	if bEnd != nil && bytes.Compare(bEnd, aStart) <= 0 {
+		return false
+	}
+	return true
+}
+
 func splitCatalogRoutes(
 	parent distribution.RouteDescriptor,
 	splitKey []byte,