Fix MGET aggregation to map individual values to keys across shards

Author: ofekshenawa

command.go

@@ -4474,7 +4474,7 @@ func (cmd *CommandsInfoCmd) readReply(rd *proto.Reader) error {
 					rawTips[k] = v
 				}
 			}
-			cmdInfo.CommandPolicy = parseCommandPolicies(rawTips)
+			cmdInfo.CommandPolicy = parseCommandPolicies(rawTips, cmdInfo.FirstKeyPos)
 
 			if err := rd.DiscardNext(); err != nil {
 				return err
@@ -4573,9 +4573,12 @@ func (c *cmdsInfoCache) Refresh() {
 const requestPolicy = "request_policy"
 const responsePolicy = "response_policy"
 
-func parseCommandPolicies(commandInfoTips map[string]string) *routing.CommandPolicy {
+func parseCommandPolicies(commandInfoTips map[string]string, firstKeyPos int8) *routing.CommandPolicy {
 	req := routing.ReqDefault
-	resp := routing.RespAllSucceeded
+	resp := routing.RespDefaultKeyless
+	if firstKeyPos > 0 {
+		resp = routing.RespDefaultHashSlot
+	}
 
 	tips := make(map[string]string, len(commandInfoTips))
 	for k, v := range commandInfoTips {

internal/routing/aggregator.go

@@ -44,9 +44,9 @@ type AggregatorResErr struct {
 func NewResponseAggregator(policy ResponsePolicy, cmdName string) ResponseAggregator {
 	switch policy {
 	case RespDefaultKeyless:
-		return &DefaultKeylessAggregator{}
+		return &DefaultKeylessAggregator{results: make([]interface{}, 0)}
 	case RespDefaultHashSlot:
-		return &DefaultKeyedAggregator{}
+		return &DefaultKeyedAggregator{results: make(map[string]interface{})}
 	case RespAllSucceeded:
 		return &AllSucceededAggregator{}
 	case RespOneSucceeded:

osscluster.go

@@ -121,6 +121,11 @@ type ClusterOptions struct {
 
 	TLSConfig *tls.Config
 
+	// DisableRoutingPolicies disables the request/response policy routing system.
+	// When disabled, all commands use the legacy routing behavior.
+	// Experimental. Will be removed when shard picker is fully implemented.
+	DisableRoutingPolicies bool
+
 	// DisableIndentity - Disable set-lib on connect.
 	//
 	// default: false
@@ -939,6 +944,29 @@ func (c *clusterState) slotRandomNode(slot int) (*clusterNode, error) {
 	return nodes[randomNodes[0]], nil
 }
 
+func (c *clusterState) slotShardPickerSlaveNode(slot int, shardPicker routing.ShardPicker) (*clusterNode, error) {
+	nodes := c.slotNodes(slot)
+	if len(nodes) == 0 {
+		return c.nodes.Random()
+	}
+
+	// nodes[0] is master, nodes[1:] are slaves
+	// First, try all slave nodes for this slot using ShardPicker order
+	slaves := nodes[1:]
+	if len(slaves) > 0 {
+		for i := 0; i < len(slaves); i++ {
+			idx := shardPicker.Next(len(slaves))
+			slave := slaves[idx]
+			if !slave.Failing() && !slave.Loading() {
+				return slave, nil
+			}
+		}
+	}
+
+	// All slaves are failing or loading - return master
+	return nodes[0], nil
+}
+
 func (c *clusterState) slotNodes(slot int) []*clusterNode {
 	i := sort.Search(len(c.slots), func(i int) bool {
 		return c.slots[i].end >= slot
@@ -1096,7 +1124,11 @@ func (c *ClusterClient) process(ctx context.Context, cmd Cmder) error {
 
 		if node == nil {
 			var err error
-			node, err = c.cmdNode(ctx, cmd.Name(), slot)
+			if !c.opt.DisableRoutingPolicies && c.opt.ShardPicker != nil {
+				node, err = c.cmdNodeWithShardPicker(ctx, cmd.Name(), slot, c.opt.ShardPicker)
+			} else {
+				node, err = c.cmdNode(ctx, cmd.Name(), slot)
+			}
 			if err != nil {
 				return err
 			}
@@ -1109,8 +1141,11 @@ func (c *ClusterClient) process(ctx context.Context, cmd Cmder) error {
 			_ = pipe.Process(ctx, cmd)
 			_, lastErr = pipe.Exec(ctx)
 		} else {
-			// Execute the command on the selected node
-			lastErr = c.routeAndRun(ctx, cmd, node)
+			if !c.opt.DisableRoutingPolicies {
+				lastErr = c.routeAndRun(ctx, cmd, node)
+			} else {
+				lastErr = node.Client.Process(ctx, cmd)
+			}
 		}
 
 		// If there is no error - we are done.
@@ -2098,13 +2133,33 @@ func (c *ClusterClient) cmdNode(
 		return nil, err
 	}
 
+	if c.opt.ReadOnly {
+		cmdInfo := c.cmdInfo(ctx, cmdName)
+		if cmdInfo != nil && cmdInfo.ReadOnly {
+			return c.slotReadOnlyNode(state, slot)
+		}
+	}
+	return state.slotMasterNode(slot)
+}
+
+func (c *ClusterClient) cmdNodeWithShardPicker(
+	ctx context.Context,
+	cmdName string,
+	slot int,
+	shardPicker routing.ShardPicker,
+) (*clusterNode, error) {
+	state, err := c.state.Get(ctx)
+	if err != nil {
+		return nil, err
+	}
+
 	// For keyless commands (slot == -1), use ShardPicker to select a shard
 	// This respects the user's configured ShardPicker policy
 	if slot == -1 {
 		if len(state.Masters) == 0 {
 			return nil, errClusterNoNodes
 		}
-		idx := c.opt.ShardPicker.Next(len(state.Masters))
+		idx := shardPicker.Next(len(state.Masters))
 		return state.Masters[idx], nil
 	}
 
@@ -2124,6 +2179,11 @@ func (c *ClusterClient) slotReadOnlyNode(state *clusterState, slot int) (*cluste
 	if c.opt.RouteRandomly {
 		return state.slotRandomNode(slot)
 	}
+
+	if c.opt.ShardPicker != nil {
+		return state.slotShardPickerSlaveNode(slot, c.opt.ShardPicker)
+	}
+
 	return state.slotSlaveNode(slot)
 }
 

osscluster_router.go

@@ -138,7 +138,7 @@ func (c *ClusterClient) executeMultiSlot(ctx context.Context, cmd Cmder, slotMap
 		go func(slot int, keys []string) {
 			defer wg.Done()
 
-			node, err := c.cmdNode(ctx, cmd.Name(), slot)
+			node, err := c.cmdNodeWithShardPicker(ctx, cmd.Name(), slot, c.opt.ShardPicker)
 			if err != nil {
 				results <- slotResult{nil, keys, err}
 				return
@@ -280,7 +280,7 @@ func (c *ClusterClient) executeCursorCommand(ctx context.Context, cmd Cmder) err
 
 	// Route based on cursor ID to maintain stickiness
 	slot := hashtag.Slot(cursorID)
-	node, err := c.cmdNode(ctx, cmd.Name(), slot)
+	node, err := c.cmdNodeWithShardPicker(ctx, cmd.Name(), slot, c.opt.ShardPicker)
 	if err != nil {
 		return err
 	}
@@ -351,12 +351,29 @@ func (c *ClusterClient) aggregateMultiSlotResults(ctx context.Context, cmd Cmder
 			firstErr = result.err
 		}
 		if result.cmd != nil && result.err == nil {
-			// For other commands, map each key to the entire result
 			value, err := ExtractCommandValue(result.cmd)
-			for _, key := range result.keys {
-				keyedResults[key] = routing.AggregatorResErr{Result: value, Err: err}
+
+			// Check if the result is a slice (e.g., from MGET)
+			if sliceValue, ok := value.([]interface{}); ok {
+				// Map each element to its corresponding key
+				for i, key := range result.keys {
+					if i < len(sliceValue) {
+						keyedResults[key] = routing.AggregatorResErr{Result: sliceValue[i], Err: err}
+					} else {
+						keyedResults[key] = routing.AggregatorResErr{Result: nil, Err: err}
+					}
+				}
+			} else {
+				// For non-slice results, map the entire result to each key
+				for _, key := range result.keys {
+					keyedResults[key] = routing.AggregatorResErr{Result: value, Err: err}
+				}
 			}
 		}
+
+		if result.err != nil {
+			firstErr = result.err
+		}
 	}
 
 	return c.aggregateKeyedValues(cmd, keyedResults, keyOrder, policy)