diff --git a/agents/drivers/rocketmq/consumers.go b/agents/drivers/rocketmq/consumers.go index 07640c2d5b..2cf1d8ef0a 100644 --- a/agents/drivers/rocketmq/consumers.go +++ b/agents/drivers/rocketmq/consumers.go @@ -273,7 +273,10 @@ func (a *rocketMQAgent) getConsumerLag(params map[string]any) (any, error) { if err != nil { return nil, err } - consumerClients := resolveConsumerClients(ctx, client, groupID, topic) + consumerClients := resolveConsumerClients(ctx, remotingConsumeStatusReader{ + client: client, + invoke: invokeRemotingWithClient, + }, groupID, topic) return buildConsumerLagResult(stats, consumerClients), nil } @@ -281,6 +284,189 @@ type consumeStatusReader interface { GetConsumeStatus(context.Context, string, string, string) (map[string]map[string]int64, error) } +type consumerStatusInvoker func( + context.Context, + string, + *remoting.RemotingCommand, +) (*remoting.RemotingCommand, error) + +type remotingConsumeStatusReader struct { + client *admin.Client + invoke consumerStatusInvoker +} + +func (r remotingConsumeStatusReader) GetConsumeStatus( + ctx context.Context, + topic string, + groupID string, + clientAddr string, +) (map[string]map[string]int64, error) { + route, err := r.client.ExamineTopicRouteInfo(ctx, topic) + if err != nil { + return nil, err + } + return readConsumerStatusFromBrokers( + ctx, masterAddressesFromRoute(route), topic, groupID, clientAddr, r.invoke, + ) +} + +func readConsumerStatusFromBrokers( + ctx context.Context, + addresses []string, + topic string, + groupID string, + clientAddr string, + invoke consumerStatusInvoker, +) (map[string]map[string]int64, error) { + if len(addresses) == 0 { + return nil, fmt.Errorf("no RocketMQ master broker found for topic %s", topic) + } + orderedAddresses := append([]string(nil), addresses...) + sort.Strings(orderedAddresses) + merged := make(map[string]map[string]int64) + successCount := 0 + var lastErr error + for _, address := range orderedAddresses { + fields := map[string]string{"topic": topic, "group": groupID} + if clientAddr != "" { + fields["clientAddr"] = clientAddr + } + response, requestErr := invoke(ctx, address, + remoting.NewRequest(remoting.InvokeBrokerToGetConsumerStatus, fields)) + if requestErr != nil { + lastErr = requestErr + continue + } + if response == nil { + lastErr = fmt.Errorf("empty consumer status response from %s", address) + continue + } + partial, decodeErr := decodeConsumerStatus(response.Body) + if decodeErr != nil { + lastErr = decodeErr + continue + } + successCount++ + for _, clientID := range sortedKeys(partial) { + if merged[clientID] == nil { + merged[clientID] = make(map[string]int64) + } + for _, queueKey := range sortedKeys(partial[clientID]) { + merged[clientID][queueKey] = partial[clientID][queueKey] + } + } + } + if successCount == 0 { + return nil, fmt.Errorf("query consumer status for group %s on all masters: %w", groupID, lastErr) + } + return merged, nil +} + +func decodeConsumerStatus(body []byte) (map[string]map[string]int64, error) { + // RocketMQ wraps assignments in GetConsumerStatusBody; admin-go v1.1.1 + // incorrectly decodes the complete response as the inner table. + var wrapper struct { + ConsumerTable map[string]map[string]int64 `json:"consumerTable"` + } + if err := json.Unmarshal(repairConsumerStatusJSON(body), &wrapper); err != nil { + return nil, fmt.Errorf("decode consumer status: %w", err) + } + if wrapper.ConsumerTable == nil { + wrapper.ConsumerTable = make(map[string]map[string]int64) + } + return wrapper.ConsumerTable, nil +} + +func repairConsumerStatusJSON(body []byte) []byte { + repaired := repairRocketMQJSON(body) + result := make([]byte, 0, len(repaired)+64) + for index := 0; index < len(repaired); { + mapStart := index + 1 + for mapStart < len(repaired) && isJSONSpace(repaired[mapStart]) { + mapStart++ + } + if repaired[index] == ':' && mapStart+1 < len(repaired) && + repaired[mapStart] == '{' && repaired[mapStart+1] == '{' { + if converted, next, ok := convertObjectKeyedInt64Map(repaired, mapStart); ok { + result = append(result, repaired[index:mapStart]...) + result = append(result, converted...) + index = next + continue + } + } + result = append(result, repaired[index]) + index++ + } + return result +} + +func convertObjectKeyedInt64Map(body []byte, start int) ([]byte, int, bool) { + result := []byte{'{'} + index := start + 1 + first := true + for index < len(body) { + for index < len(body) && isJSONSpace(body[index]) { + index++ + } + if index >= len(body) { + return nil, start, false + } + if body[index] == '}' { + return append(result, '}'), index + 1, true + } + if body[index] == ',' { + index++ + continue + } + if body[index] != '{' { + return nil, start, false + } + keyEnd := matchingBrace(body, index) + if keyEnd < 0 { + return nil, start, false + } + key := body[index : keyEnd+1] + index = keyEnd + 1 + for index < len(body) && isJSONSpace(body[index]) { + index++ + } + if index >= len(body) || body[index] != ':' { + return nil, start, false + } + index++ + for index < len(body) && isJSONSpace(body[index]) { + index++ + } + valueStart := index + if index < len(body) && body[index] == '-' { + index++ + } + digitStart := index + for index < len(body) && body[index] >= '0' && body[index] <= '9' { + index++ + } + if index == digitStart { + return nil, start, false + } + valueEnd := index + for index < len(body) && isJSONSpace(body[index]) { + index++ + } + if index >= len(body) || (body[index] != ',' && body[index] != '}') { + return nil, start, false + } + if !first { + result = append(result, ',') + } + first = false + result = append(result, '"') + result = append(result, escapeJSONString(key)...) + result = append(result, '"', ':') + result = append(result, body[valueStart:valueEnd]...) + } + return nil, start, false +} + func resolveConsumerClients( ctx context.Context, reader consumeStatusReader, diff --git a/agents/drivers/rocketmq/consumers_test.go b/agents/drivers/rocketmq/consumers_test.go index 56ba320c20..a62427b59a 100644 --- a/agents/drivers/rocketmq/consumers_test.go +++ b/agents/drivers/rocketmq/consumers_test.go @@ -9,6 +9,7 @@ import ( "testing" admin "github.com/amigoer/rocketmq-admin-go" + "github.com/amigoer/rocketmq-admin-go/protocol/remoting" ) type fakeConsumeStatusReader struct { @@ -68,6 +69,112 @@ func TestDecodeConsumeStatsRepairsObjectKeys(t *testing.T) { } } +func TestDecodeConsumerStatusUnwrapsRocketMQ48Body(t *testing.T) { + body := []byte(`{ + "messageQueueTable": {}, + "consumerTable": { + "client-b": {{"brokerName":"broker-a","queueId":0,"topic":"Orders"}:90}, + "client-a": { + "MessageQueue [topic=Orders, brokerName=broker-a, queueId=1]": 50 + } + } + }`) + + status, err := decodeConsumerStatus(body) + if err != nil { + t.Fatal(err) + } + want := map[string]map[string]int64{ + "client-b": { + `{"brokerName":"broker-a","queueId":0,"topic":"Orders"}`: 90, + }, + "client-a": { + `MessageQueue [topic=Orders, brokerName=broker-a, queueId=1]`: 50, + }, + } + if !reflect.DeepEqual(status, want) { + t.Fatalf("consumer status = %#v, want %#v", status, want) + } +} + +func TestDecodeConsumerStatusHandlesEmptyAndMalformedBodies(t *testing.T) { + status, err := decodeConsumerStatus([]byte(`{"messageQueueTable":{},"consumerTable":{}}`)) + if err != nil { + t.Fatal(err) + } + if len(status) != 0 { + t.Fatalf("empty consumer table decoded as %#v", status) + } + if _, err := decodeConsumerStatus([]byte(`{"consumerTable":[]}`)); err == nil { + t.Fatal("expected malformed consumer table to fail") + } +} + +func TestReadConsumerStatusFromBrokersUsesBoundedRequestsAndMerges(t *testing.T) { + responses := map[string][]byte{ + "broker-a:10911": []byte(`{"consumerTable":{"client-a":{"{\"topic\":\"Orders\",\"brokerName\":\"broker-a\",\"queueId\":0}":90,"{\"topic\":\"Orders\",\"brokerName\":\"broker-a\",\"queueId\":1}":50}}}`), + "broker-b:10911": []byte(`{"consumerTable":{"client-a":{"MessageQueue [topic=Orders, brokerName=broker-b, queueId=0]":20}}}`), + } + calls := make([]string, 0, len(responses)) + invoke := func(_ context.Context, address string, command *remoting.RemotingCommand) (*remoting.RemotingCommand, error) { + calls = append(calls, address) + if command.Code != remoting.InvokeBrokerToGetConsumerStatus { + t.Fatalf("request code = %d, want %d", command.Code, remoting.InvokeBrokerToGetConsumerStatus) + } + if command.ExtFields["topic"] != "Orders" || command.ExtFields["group"] != "GID_Orders" { + t.Fatalf("unexpected request fields: %#v", command.ExtFields) + } + if _, exists := command.ExtFields["clientAddr"]; exists { + t.Fatalf("empty client address sent on wire: %#v", command.ExtFields) + } + return &remoting.RemotingCommand{Code: remoting.Success, Body: responses[address]}, nil + } + + status, err := readConsumerStatusFromBrokers( + context.Background(), []string{"broker-b:10911", "broker-a:10911"}, + "Orders", "GID_Orders", "", invoke, + ) + if err != nil { + t.Fatal(err) + } + if !reflect.DeepEqual(calls, []string{"broker-a:10911", "broker-b:10911"}) { + t.Fatalf("broker calls = %#v", calls) + } + if len(status["client-a"]) != 3 { + t.Fatalf("merged status = %#v", status) + } +} + +func TestReadConsumerStatusFromBrokersKeepsPartialResults(t *testing.T) { + invoke := func(_ context.Context, address string, _ *remoting.RemotingCommand) (*remoting.RemotingCommand, error) { + if address == "broker-b:10911" { + return nil, errors.New("consumer offline") + } + return &remoting.RemotingCommand{Code: remoting.Success, Body: []byte( + `{"consumerTable":{"client-a":{"{\"topic\":\"Orders\",\"brokerName\":\"broker-a\",\"queueId\":0}":90}}}`, + )}, nil + } + + status, err := readConsumerStatusFromBrokers( + context.Background(), []string{"broker-a:10911", "broker-b:10911"}, + "Orders", "GID_Orders", "", invoke, + ) + if err != nil { + t.Fatal(err) + } + if len(status["client-a"]) != 1 { + t.Fatalf("partial status = %#v", status) + } + + _, err = readConsumerStatusFromBrokers( + context.Background(), []string{"broker-b:10911"}, + "Orders", "GID_Orders", "", invoke, + ) + if err == nil { + t.Fatal("expected all-broker request failure") + } +} + func TestResolveConsumerClientsUsesOneBestEffortLookup(t *testing.T) { reader := &fakeConsumeStatusReader{status: map[string]map[string]int64{ "client-z": {