refactor: [shard-distributor]Handle context.Cancelled errors

service/sharddistributor/leader/process/processor.go

@@ -215,6 +215,9 @@ func (p *namespaceProcessor) runRebalancingLoop(ctx context.Context) {
 			err = p.rebalanceShards(ctx)
 		}
 		if err != nil {
+			if store.IsContextCancellation(err) {
+				return
+			}
 			p.logger.Error("rebalance failed", tag.Error(err))
 		}
 	}
@@ -234,6 +237,9 @@ func (p *namespaceProcessor) runShardStatsCleanupLoop(ctx context.Context) {
 			p.logger.Info("Periodic shard stats cleanup triggered.")
 			namespaceState, err := p.shardStore.GetState(ctx, p.namespaceCfg.Name)
 			if err != nil {
+				if store.IsContextCancellation(err) {
+					return
+				}
 				p.logger.Error("Failed to get state for shard stats cleanup", tag.Error(err))
 				continue
 			}
@@ -243,6 +249,9 @@ func (p *namespaceProcessor) runShardStatsCleanupLoop(ctx context.Context) {
 				continue
 			}
 			if err := p.shardStore.DeleteShardStats(ctx, p.namespaceCfg.Name, staleShardStats, p.election.Guard()); err != nil {
+				if store.IsContextCancellation(err) {
+					return
+				}
 				p.logger.Error("Failed to delete stale shard stats", tag.Error(err))
 			}
 		}
@@ -340,6 +349,9 @@ func (p *namespaceProcessor) rebalanceShards(ctx context.Context) (err error) {
 func (p *namespaceProcessor) rebalanceShardsImpl(ctx context.Context, metricsLoopScope metrics.Scope) (err error) {
 	namespaceState, err := p.shardStore.GetState(ctx, p.namespaceCfg.Name)
 	if err != nil {
+		if store.IsContextCancellation(err) {
+			return err
+		}
 		return fmt.Errorf("get state: %w", err)
 	}
 
@@ -386,6 +398,9 @@ func (p *namespaceProcessor) rebalanceShardsImpl(ctx context.Context, metricsLoo
 		ExecutorsToDelete: staleExecutors,
 	}, p.election.Guard())
 	if err != nil {
+		if store.IsContextCancellation(err) {
+			return err
+		}
 		return fmt.Errorf("assign shards: %w", err)
 	}
 

service/sharddistributor/store/etcd/executorstore/etcdstore.go

@@ -233,8 +233,14 @@ func (s *executorStoreImpl) GetState(ctx context.Context, namespace string) (*st
 	executorPrefix := etcdkeys.BuildExecutorsPrefix(s.prefix, namespace)
 	resp, err := s.client.Get(ctx, executorPrefix, clientv3.WithPrefix())
 	if err != nil {
+		if store.IsContextCancellation(err) {
+			return nil, ctx.Err()
+		}
 		return nil, fmt.Errorf("get executor data: %w", err)
 	}
+	if ctxErr := ctx.Err(); store.IsContextCancellation(ctxErr) {
+		return nil, ctxErr
+	}
 
 	for _, kv := range resp.Kvs {
 		key := string(kv.Key)
@@ -276,8 +282,14 @@ func (s *executorStoreImpl) GetState(ctx context.Context, namespace string) (*st
 	shardsPrefix := etcdkeys.BuildShardsPrefix(s.prefix, namespace)
 	shardResp, err := s.client.Get(ctx, shardsPrefix, clientv3.WithPrefix())
 	if err != nil {
+		if store.IsContextCancellation(err) {
+			return nil, ctx.Err()
+		}
 		return nil, fmt.Errorf("get shard data: %w", err)
 	}
+	if ctxErr := ctx.Err(); store.IsContextCancellation(ctxErr) {
+		return nil, ctxErr
+	}
 	for _, kv := range shardResp.Kvs {
 		shardID, shardKeyType, err := etcdkeys.ParseShardKey(s.prefix, namespace, string(kv.Key))
 		if err != nil {

service/sharddistributor/store/store.go

@@ -2,6 +2,7 @@ package store
 
 import (
 	"context"
+	"errors"
 	"fmt"
 )
 
@@ -48,6 +49,16 @@ func NopGuard() GuardFunc {
 	}
 }
 
+// IsContextCancellation reports whether the provided error indicates the caller's context
+// has been cancelled or its deadline has been exceeded.
+func IsContextCancellation(err error) bool {
+	if err == nil {
+		return false
+	}
+	return errors.Is(err, context.Canceled) ||
+		errors.Is(err, context.DeadlineExceeded)
+}
+
 // AssignShardsRequest is a request to assign shards to executors, and remove unused shards.
 type AssignShardsRequest struct {
 	// NewState is the new state of the namespace, containing the new assignments of shards to executors.