feat(unified-storage): implement scatter-gather queries with replica failover
Phase 3 of the sub-index sharding implementation: - Add parallelSearchWithFailover() that fans out queries to all sub-indexes - Implement replica failover - tries each replica in order until success - Add mergeResults() with deduplication, sorting, pagination - Handle partial failures with HTTP 206 status when some shards fail - Add comprehensive unit tests for merge, dedup, sort, and facet operations The scatter-gather pattern enables distributed search across sub-indexes while maintaining the ≤250ms SLO target through parallel execution and failover.
This commit is contained in:
Rafael Paulovic
@@ -1,11 +1,14 @@
|
||||
package resource
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"math/rand"
|
||||
"slices"
|
||||
"sort"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -15,6 +18,7 @@ import (
|
||||
"github.com/grafana/dskit/services"
|
||||
userutils "github.com/grafana/dskit/user"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/trace"
|
||||
"google.golang.org/grpc"
|
||||
"google.golang.org/grpc/health/grpc_health_v1"
|
||||
@@ -35,10 +39,12 @@ func ProvideSearchDistributorServer(cfg *setting.Cfg, features featuremgmt.Featu
|
||||
}
|
||||
|
||||
distributorServer := &distributorServer{
|
||||
log: log.New("index-server-distributor"),
|
||||
ring: ring,
|
||||
clientPool: ringClientPool,
|
||||
tracing: tracer,
|
||||
log: log.New("index-server-distributor"),
|
||||
ring: ring,
|
||||
clientPool: ringClientPool,
|
||||
tracing: tracer,
|
||||
subIndexesPerNamespace: cfg.SubIndexesPerNamespace,
|
||||
replicationFactor: cfg.SearchRingReplicationFactor,
|
||||
}
|
||||
|
||||
healthService, err := ProvideHealthService(distributorServer)
|
||||
@@ -83,10 +89,12 @@ const RingHeartbeatTimeout = time.Minute
|
||||
const RingNumTokens = 128
|
||||
|
||||
type distributorServer struct {
|
||||
clientPool *ringclient.Pool
|
||||
ring *ring.Ring
|
||||
log log.Logger
|
||||
tracing trace.Tracer
|
||||
clientPool *ringclient.Pool
|
||||
ring *ring.Ring
|
||||
log log.Logger
|
||||
tracing trace.Tracer
|
||||
subIndexesPerNamespace int // Number of sub-indexes per namespace (0 = disabled)
|
||||
replicationFactor int // Ring replication factor for replica failover
|
||||
}
|
||||
|
||||
var (
|
||||
@@ -99,12 +107,33 @@ var (
|
||||
func (ds *distributorServer) Search(ctx context.Context, r *resourcepb.ResourceSearchRequest) (*resourcepb.ResourceSearchResponse, error) {
|
||||
ctx, span := ds.tracing.Start(ctx, "distributor.Search")
|
||||
defer span.End()
|
||||
ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Options.Key.Namespace, "Search")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
// If sub-index sharding is not enabled, use the existing single-node routing
|
||||
if ds.subIndexesPerNamespace <= 0 {
|
||||
ctx, client, err := ds.getClientToDistributeRequest(ctx, r.Options.Key.Namespace, "Search")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return client.Search(ctx, r)
|
||||
}
|
||||
|
||||
return client.Search(ctx, r)
|
||||
// Scatter-gather search across all sub-indexes
|
||||
nsr := NamespacedResource{
|
||||
Namespace: r.Options.Key.Namespace,
|
||||
Group: r.Options.Key.Group,
|
||||
Resource: r.Options.Key.Resource,
|
||||
}
|
||||
|
||||
span.SetAttributes(
|
||||
attribute.String("namespace", nsr.Namespace),
|
||||
attribute.String("group", nsr.Group),
|
||||
attribute.String("resource", nsr.Resource),
|
||||
attribute.Int("sub_indexes", ds.subIndexesPerNamespace),
|
||||
)
|
||||
|
||||
subIndexes := ds.getSubIndexesForNamespace(nsr)
|
||||
results := ds.parallelSearchWithFailover(ctx, subIndexes, r)
|
||||
return ds.mergeResults(ctx, results, r)
|
||||
}
|
||||
|
||||
func (ds *distributorServer) GetStats(ctx context.Context, r *resourcepb.ResourceStatsRequest) (*resourcepb.ResourceStatsResponse, error) {
|
||||
@@ -276,3 +305,383 @@ func (ds *distributorServer) IsHealthy(ctx context.Context, r *resourcepb.Health
|
||||
|
||||
return &resourcepb.HealthCheckResponse{Status: resourcepb.HealthCheckResponse_NOT_SERVING}, nil
|
||||
}
|
||||
|
||||
// --- Scatter-Gather Query Implementation with Replica Failover ---
|
||||
|
||||
// subIndexSearchResult holds the result from searching a single sub-index.
|
||||
type subIndexSearchResult struct {
|
||||
subIndexID int
|
||||
response *resourcepb.ResourceSearchResponse
|
||||
err error
|
||||
partialFailure bool // true if all replicas for this sub-index failed
|
||||
}
|
||||
|
||||
// getSubIndexesForNamespace returns all sub-index keys for a NamespacedResource.
|
||||
// This is used for scatter-gather queries that need to query all sub-indexes.
|
||||
func (ds *distributorServer) getSubIndexesForNamespace(nsr NamespacedResource) []SubIndexKey {
|
||||
count := ds.subIndexesPerNamespace
|
||||
if count <= 0 {
|
||||
count = 1
|
||||
}
|
||||
keys := make([]SubIndexKey, count)
|
||||
for i := 0; i < count; i++ {
|
||||
keys[i] = SubIndexKey{
|
||||
NamespacedResource: nsr,
|
||||
SubIndexID: i,
|
||||
}
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
// getReplicasForSubIndex returns the ordered list of replicas (instances) that own
|
||||
// the given sub-index. Replicas are ordered by preference from the ring.
|
||||
// Uses consistent hashing including the sub-index ID to determine ownership.
|
||||
func (ds *distributorServer) getReplicasForSubIndex(subIndex SubIndexKey) ([]ring.InstanceDesc, error) {
|
||||
ringHasher := fnv.New32a()
|
||||
// Include sub-index ID in hash to distribute sub-indexes across nodes
|
||||
_, err := ringHasher.Write([]byte(fmt.Sprintf("%s/%d", subIndex.Namespace, subIndex.SubIndexID)))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error hashing sub-index key: %w", err)
|
||||
}
|
||||
|
||||
replicationFactor := ds.replicationFactor
|
||||
if replicationFactor <= 0 {
|
||||
replicationFactor = ds.ring.ReplicationFactor()
|
||||
}
|
||||
|
||||
rs, err := ds.ring.GetWithOptions(ringHasher.Sum32(), searchRingRead, ring.WithReplicationFactor(replicationFactor))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("error getting replication set from ring for sub-index %s: %w", subIndex.String(), err)
|
||||
}
|
||||
|
||||
return rs.Instances, nil
|
||||
}
|
||||
|
||||
// parallelSearchWithFailover executes search queries across all sub-indexes in parallel.
|
||||
// For each sub-index, it tries replicas in order until one succeeds.
|
||||
// Returns results from all sub-indexes (some may be marked as partial failures).
|
||||
func (ds *distributorServer) parallelSearchWithFailover(ctx context.Context, subIndexes []SubIndexKey, r *resourcepb.ResourceSearchRequest) []*subIndexSearchResult {
|
||||
ctx, span := ds.tracing.Start(ctx, "distributor.parallelSearchWithFailover")
|
||||
defer span.End()
|
||||
|
||||
results := make([]*subIndexSearchResult, len(subIndexes))
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for i, subIdx := range subIndexes {
|
||||
wg.Add(1)
|
||||
go func(idx int, key SubIndexKey) {
|
||||
defer wg.Done()
|
||||
results[idx] = ds.searchSubIndexWithFailover(ctx, key, r)
|
||||
}(i, subIdx)
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
// Count partial failures for logging
|
||||
partialFailures := 0
|
||||
for _, result := range results {
|
||||
if result.partialFailure {
|
||||
partialFailures++
|
||||
}
|
||||
}
|
||||
span.SetAttributes(attribute.Int("partial_failures", partialFailures))
|
||||
|
||||
return results
|
||||
}
|
||||
|
||||
// searchSubIndexWithFailover searches a single sub-index, trying each replica in order
|
||||
// until one succeeds. If all replicas fail, marks the result as a partial failure.
|
||||
func (ds *distributorServer) searchSubIndexWithFailover(ctx context.Context, subIndex SubIndexKey, r *resourcepb.ResourceSearchRequest) *subIndexSearchResult {
|
||||
result := &subIndexSearchResult{
|
||||
subIndexID: subIndex.SubIndexID,
|
||||
}
|
||||
|
||||
// Get ordered list of replicas for this sub-index
|
||||
replicas, err := ds.getReplicasForSubIndex(subIndex)
|
||||
if err != nil {
|
||||
ds.log.Warn("failed to get replicas for sub-index", "subIndex", subIndex.String(), "error", err)
|
||||
result.err = err
|
||||
result.partialFailure = true
|
||||
return result
|
||||
}
|
||||
|
||||
if len(replicas) == 0 {
|
||||
ds.log.Warn("no replicas available for sub-index", "subIndex", subIndex.String())
|
||||
result.err = fmt.Errorf("no replicas available for sub-index %s", subIndex.String())
|
||||
result.partialFailure = true
|
||||
return result
|
||||
}
|
||||
|
||||
// Prepare context with metadata
|
||||
md, ok := metadata.FromIncomingContext(ctx)
|
||||
if !ok {
|
||||
md = make(metadata.MD)
|
||||
}
|
||||
rCtx := userutils.InjectOrgID(metadata.NewOutgoingContext(ctx, md), subIndex.Namespace)
|
||||
|
||||
// Try each replica in order until success
|
||||
var lastErr error
|
||||
for replicaIdx, replica := range replicas {
|
||||
client, err := ds.clientPool.GetClientForInstance(replica)
|
||||
if err != nil {
|
||||
ds.log.Debug("failed to get client for replica",
|
||||
"subIndex", subIndex.String(),
|
||||
"replica", replica.Id,
|
||||
"replicaIdx", replicaIdx,
|
||||
"error", err)
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
|
||||
resp, err := client.(*RingClient).Client.Search(rCtx, r)
|
||||
if err == nil && resp.Error == nil {
|
||||
// Success
|
||||
result.response = resp
|
||||
return result
|
||||
}
|
||||
|
||||
// Log failover event
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
ds.log.Warn("search failed on replica, failing over to next",
|
||||
"subIndex", subIndex.String(),
|
||||
"replica", replica.Id,
|
||||
"replicaIdx", replicaIdx,
|
||||
"remainingReplicas", len(replicas)-replicaIdx-1,
|
||||
"error", err)
|
||||
} else if resp.Error != nil {
|
||||
lastErr = fmt.Errorf("search error: %s", resp.Error.Message)
|
||||
ds.log.Warn("search returned error on replica, failing over to next",
|
||||
"subIndex", subIndex.String(),
|
||||
"replica", replica.Id,
|
||||
"replicaIdx", replicaIdx,
|
||||
"remainingReplicas", len(replicas)-replicaIdx-1,
|
||||
"errorMessage", resp.Error.Message)
|
||||
}
|
||||
}
|
||||
|
||||
// All replicas failed - mark as partial failure
|
||||
ds.log.Error("all replicas failed for sub-index",
|
||||
"subIndex", subIndex.String(),
|
||||
"totalReplicas", len(replicas),
|
||||
"lastError", lastErr)
|
||||
result.err = lastErr
|
||||
result.partialFailure = true
|
||||
return result
|
||||
}
|
||||
|
||||
// mergeResults combines results from all sub-indexes into a single response.
|
||||
// It handles:
|
||||
// - Result deduplication by resource key
|
||||
// - Sort merging (merge-sort for sorted results)
|
||||
// - Pagination across shards
|
||||
// - Tracking and reporting partial failures
|
||||
func (ds *distributorServer) mergeResults(ctx context.Context, results []*subIndexSearchResult, req *resourcepb.ResourceSearchRequest) (*resourcepb.ResourceSearchResponse, error) {
|
||||
ctx, span := ds.tracing.Start(ctx, "distributor.mergeResults")
|
||||
defer span.End()
|
||||
|
||||
// Collect all rows and track partial failures
|
||||
var allRows []*resourcepb.ResourceTableRow
|
||||
var columns []*resourcepb.ResourceTableColumnDefinition
|
||||
var totalHits int64
|
||||
var totalQueryCost float64
|
||||
var maxScore float64
|
||||
partialFailures := 0
|
||||
failedSubIndexes := []int{}
|
||||
facets := make(map[string]*resourcepb.ResourceSearchResponse_Facet)
|
||||
|
||||
for _, result := range results {
|
||||
if result.partialFailure {
|
||||
partialFailures++
|
||||
failedSubIndexes = append(failedSubIndexes, result.subIndexID)
|
||||
continue
|
||||
}
|
||||
|
||||
if result.response == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
resp := result.response
|
||||
totalHits += resp.TotalHits
|
||||
totalQueryCost += resp.QueryCost
|
||||
if resp.MaxScore > maxScore {
|
||||
maxScore = resp.MaxScore
|
||||
}
|
||||
|
||||
// Use columns from first valid response
|
||||
if columns == nil && resp.Results != nil {
|
||||
columns = resp.Results.Columns
|
||||
}
|
||||
|
||||
// Collect rows
|
||||
if resp.Results != nil && len(resp.Results.Rows) > 0 {
|
||||
allRows = append(allRows, resp.Results.Rows...)
|
||||
}
|
||||
|
||||
// Merge facets
|
||||
for k, v := range resp.Facet {
|
||||
if existing, ok := facets[k]; ok {
|
||||
mergeFacets(existing, v)
|
||||
} else {
|
||||
facets[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
span.SetAttributes(
|
||||
attribute.Int("total_rows_before_dedup", len(allRows)),
|
||||
attribute.Int("partial_failures", partialFailures),
|
||||
)
|
||||
|
||||
// Deduplicate rows by resource key
|
||||
allRows = deduplicateRows(allRows)
|
||||
|
||||
span.SetAttributes(attribute.Int("total_rows_after_dedup", len(allRows)))
|
||||
|
||||
// Sort rows if sort criteria provided
|
||||
if len(req.SortBy) > 0 && columns != nil {
|
||||
sortRows(allRows, columns, req.SortBy)
|
||||
}
|
||||
|
||||
// Apply pagination
|
||||
offset := int(req.Offset)
|
||||
limit := int(req.Limit)
|
||||
if limit <= 0 {
|
||||
limit = 100 // default limit
|
||||
}
|
||||
|
||||
var paginatedRows []*resourcepb.ResourceTableRow
|
||||
if offset < len(allRows) {
|
||||
end := offset + limit
|
||||
if end > len(allRows) {
|
||||
end = len(allRows)
|
||||
}
|
||||
paginatedRows = allRows[offset:end]
|
||||
}
|
||||
|
||||
// Build response
|
||||
response := &resourcepb.ResourceSearchResponse{
|
||||
TotalHits: totalHits,
|
||||
QueryCost: totalQueryCost,
|
||||
MaxScore: maxScore,
|
||||
Results: &resourcepb.ResourceTable{
|
||||
Columns: columns,
|
||||
Rows: paginatedRows,
|
||||
},
|
||||
Facet: facets,
|
||||
}
|
||||
|
||||
// Report partial failures if any
|
||||
if partialFailures > 0 {
|
||||
response.Error = &resourcepb.ErrorResult{
|
||||
Code: 206, // Partial Content
|
||||
Message: fmt.Sprintf("partial results: %d of %d sub-indexes failed (sub-indexes: %v)", partialFailures, len(results), failedSubIndexes),
|
||||
}
|
||||
ds.log.Warn("search returned partial results",
|
||||
"namespace", req.Options.Key.Namespace,
|
||||
"failedSubIndexes", partialFailures,
|
||||
"totalSubIndexes", len(results))
|
||||
}
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
// deduplicateRows removes duplicate rows based on their resource key.
|
||||
// If duplicates exist, keeps the first occurrence.
|
||||
func deduplicateRows(rows []*resourcepb.ResourceTableRow) []*resourcepb.ResourceTableRow {
|
||||
if len(rows) == 0 {
|
||||
return rows
|
||||
}
|
||||
|
||||
seen := make(map[string]bool, len(rows))
|
||||
result := make([]*resourcepb.ResourceTableRow, 0, len(rows))
|
||||
|
||||
for _, row := range rows {
|
||||
if row.Key == nil {
|
||||
continue
|
||||
}
|
||||
key := SearchID(row.Key)
|
||||
if !seen[key] {
|
||||
seen[key] = true
|
||||
result = append(result, row)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
// sortRows sorts rows based on the sort criteria.
|
||||
// Uses stable sort to maintain relative ordering of equal elements.
|
||||
func sortRows(rows []*resourcepb.ResourceTableRow, columns []*resourcepb.ResourceTableColumnDefinition, sortBy []*resourcepb.ResourceSearchRequest_Sort) {
|
||||
if len(rows) == 0 || len(sortBy) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
// Build column index map
|
||||
columnIndex := make(map[string]int)
|
||||
for i, col := range columns {
|
||||
columnIndex[col.Name] = i
|
||||
}
|
||||
|
||||
sort.SliceStable(rows, func(i, j int) bool {
|
||||
for _, s := range sortBy {
|
||||
colIdx, ok := columnIndex[s.Field]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
|
||||
// Get cell values
|
||||
var valI, valJ []byte
|
||||
if colIdx < len(rows[i].Cells) {
|
||||
valI = rows[i].Cells[colIdx]
|
||||
}
|
||||
if colIdx < len(rows[j].Cells) {
|
||||
valJ = rows[j].Cells[colIdx]
|
||||
}
|
||||
|
||||
// Compare byte slices
|
||||
cmpResult := slices.Compare(valI, valJ)
|
||||
if cmpResult == 0 {
|
||||
continue // Values are equal, check next sort field
|
||||
}
|
||||
|
||||
// Apply descending order if needed
|
||||
if s.Desc {
|
||||
return cmpResult > 0
|
||||
}
|
||||
return cmpResult < 0
|
||||
}
|
||||
return false // All sort fields are equal
|
||||
})
|
||||
}
|
||||
|
||||
// mergeFacets merges facet data from source into target.
|
||||
func mergeFacets(target, source *resourcepb.ResourceSearchResponse_Facet) {
|
||||
if source == nil {
|
||||
return
|
||||
}
|
||||
|
||||
target.Total += source.Total
|
||||
target.Missing += source.Missing
|
||||
|
||||
// Merge term facets
|
||||
termMap := make(map[string]int64)
|
||||
for _, t := range target.Terms {
|
||||
termMap[t.Term] = t.Count
|
||||
}
|
||||
for _, t := range source.Terms {
|
||||
termMap[t.Term] += t.Count
|
||||
}
|
||||
|
||||
// Rebuild term slice sorted by count (descending)
|
||||
target.Terms = make([]*resourcepb.ResourceSearchResponse_TermFacet, 0, len(termMap))
|
||||
for term, count := range termMap {
|
||||
target.Terms = append(target.Terms, &resourcepb.ResourceSearchResponse_TermFacet{
|
||||
Term: term,
|
||||
Count: count,
|
||||
})
|
||||
}
|
||||
slices.SortFunc(target.Terms, func(a, b *resourcepb.ResourceSearchResponse_TermFacet) int {
|
||||
return cmp.Compare(b.Count, a.Count) // Descending order
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,390 @@
|
||||
package resource
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.opentelemetry.io/otel/trace/noop"
|
||||
|
||||
"github.com/grafana/grafana/pkg/infra/log"
|
||||
"github.com/grafana/grafana/pkg/storage/unified/resourcepb"
|
||||
)
|
||||
|
||||
func TestGetSubIndexesForNamespace(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
subIndexCount int
|
||||
expectedCount int
|
||||
expectedIDs []int
|
||||
}{
|
||||
{
|
||||
name: "zero sub-indexes defaults to 1",
|
||||
subIndexCount: 0,
|
||||
expectedCount: 1,
|
||||
expectedIDs: []int{0},
|
||||
},
|
||||
{
|
||||
name: "negative sub-indexes defaults to 1",
|
||||
subIndexCount: -1,
|
||||
expectedCount: 1,
|
||||
expectedIDs: []int{0},
|
||||
},
|
||||
{
|
||||
name: "4 sub-indexes",
|
||||
subIndexCount: 4,
|
||||
expectedCount: 4,
|
||||
expectedIDs: []int{0, 1, 2, 3},
|
||||
},
|
||||
{
|
||||
name: "64 sub-indexes",
|
||||
subIndexCount: 64,
|
||||
expectedCount: 64,
|
||||
expectedIDs: nil, // Don't check all 64
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
ds := &distributorServer{
|
||||
subIndexesPerNamespace: tt.subIndexCount,
|
||||
}
|
||||
|
||||
nsr := NamespacedResource{
|
||||
Namespace: "org-1",
|
||||
Group: "dashboard.grafana.app",
|
||||
Resource: "dashboards",
|
||||
}
|
||||
|
||||
result := ds.getSubIndexesForNamespace(nsr)
|
||||
assert.Len(t, result, tt.expectedCount)
|
||||
|
||||
// Check that all keys have the correct NSR
|
||||
for i, key := range result {
|
||||
assert.Equal(t, nsr, key.NamespacedResource)
|
||||
assert.Equal(t, i, key.SubIndexID)
|
||||
}
|
||||
|
||||
// Check specific IDs if provided
|
||||
if tt.expectedIDs != nil {
|
||||
for i, expectedID := range tt.expectedIDs {
|
||||
assert.Equal(t, expectedID, result[i].SubIndexID)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeduplicateRows(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
input []*resourcepb.ResourceTableRow
|
||||
expected int // expected number of rows after dedup
|
||||
}{
|
||||
{
|
||||
name: "empty input",
|
||||
input: []*resourcepb.ResourceTableRow{},
|
||||
expected: 0,
|
||||
},
|
||||
{
|
||||
name: "no duplicates",
|
||||
input: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "c"}},
|
||||
},
|
||||
expected: 3,
|
||||
},
|
||||
{
|
||||
name: "with duplicates",
|
||||
input: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}}, // duplicate
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "c"}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}}, // duplicate
|
||||
},
|
||||
expected: 3,
|
||||
},
|
||||
{
|
||||
name: "rows with nil keys are skipped",
|
||||
input: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a"}},
|
||||
{Key: nil},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b"}},
|
||||
},
|
||||
expected: 2,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := deduplicateRows(tt.input)
|
||||
assert.Len(t, result, tt.expected)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSortRows(t *testing.T) {
|
||||
columns := []*resourcepb.ResourceTableColumnDefinition{
|
||||
{Name: "title"},
|
||||
{Name: "created"},
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
rows []*resourcepb.ResourceTableRow
|
||||
sortBy []*resourcepb.ResourceSearchRequest_Sort
|
||||
expectedOrder []string // expected order of first cell values
|
||||
}{
|
||||
{
|
||||
name: "sort ascending by title",
|
||||
rows: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Name: "c"}, Cells: [][]byte{[]byte("charlie"), nil}},
|
||||
{Key: &resourcepb.ResourceKey{Name: "a"}, Cells: [][]byte{[]byte("alpha"), nil}},
|
||||
{Key: &resourcepb.ResourceKey{Name: "b"}, Cells: [][]byte{[]byte("bravo"), nil}},
|
||||
},
|
||||
sortBy: []*resourcepb.ResourceSearchRequest_Sort{
|
||||
{Field: "title", Desc: false},
|
||||
},
|
||||
expectedOrder: []string{"alpha", "bravo", "charlie"},
|
||||
},
|
||||
{
|
||||
name: "sort descending by title",
|
||||
rows: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Name: "a"}, Cells: [][]byte{[]byte("alpha"), nil}},
|
||||
{Key: &resourcepb.ResourceKey{Name: "b"}, Cells: [][]byte{[]byte("bravo"), nil}},
|
||||
{Key: &resourcepb.ResourceKey{Name: "c"}, Cells: [][]byte{[]byte("charlie"), nil}},
|
||||
},
|
||||
sortBy: []*resourcepb.ResourceSearchRequest_Sort{
|
||||
{Field: "title", Desc: true},
|
||||
},
|
||||
expectedOrder: []string{"charlie", "bravo", "alpha"},
|
||||
},
|
||||
{
|
||||
name: "empty rows",
|
||||
rows: []*resourcepb.ResourceTableRow{},
|
||||
sortBy: []*resourcepb.ResourceSearchRequest_Sort{{Field: "title"}},
|
||||
expectedOrder: []string{},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
sortRows(tt.rows, columns, tt.sortBy)
|
||||
|
||||
for i, expected := range tt.expectedOrder {
|
||||
if i < len(tt.rows) {
|
||||
assert.Equal(t, expected, string(tt.rows[i].Cells[0]))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeFacets(t *testing.T) {
|
||||
t.Run("merge term facets", func(t *testing.T) {
|
||||
target := &resourcepb.ResourceSearchResponse_Facet{
|
||||
Field: "tags",
|
||||
Total: 100,
|
||||
Missing: 5,
|
||||
Terms: []*resourcepb.ResourceSearchResponse_TermFacet{
|
||||
{Term: "production", Count: 50},
|
||||
{Term: "staging", Count: 30},
|
||||
},
|
||||
}
|
||||
|
||||
source := &resourcepb.ResourceSearchResponse_Facet{
|
||||
Field: "tags",
|
||||
Total: 80,
|
||||
Missing: 3,
|
||||
Terms: []*resourcepb.ResourceSearchResponse_TermFacet{
|
||||
{Term: "production", Count: 40},
|
||||
{Term: "development", Count: 25},
|
||||
},
|
||||
}
|
||||
|
||||
mergeFacets(target, source)
|
||||
|
||||
assert.Equal(t, int64(180), target.Total)
|
||||
assert.Equal(t, int64(8), target.Missing)
|
||||
assert.Len(t, target.Terms, 3)
|
||||
|
||||
// Terms should be sorted by count descending
|
||||
termCounts := make(map[string]int64)
|
||||
for _, term := range target.Terms {
|
||||
termCounts[term.Term] = term.Count
|
||||
}
|
||||
assert.Equal(t, int64(90), termCounts["production"])
|
||||
assert.Equal(t, int64(30), termCounts["staging"])
|
||||
assert.Equal(t, int64(25), termCounts["development"])
|
||||
})
|
||||
|
||||
t.Run("merge nil source", func(t *testing.T) {
|
||||
target := &resourcepb.ResourceSearchResponse_Facet{
|
||||
Total: 100,
|
||||
}
|
||||
mergeFacets(target, nil)
|
||||
assert.Equal(t, int64(100), target.Total)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSubIndexSearchResult(t *testing.T) {
|
||||
t.Run("successful result", func(t *testing.T) {
|
||||
result := &subIndexSearchResult{
|
||||
subIndexID: 5,
|
||||
response: &resourcepb.ResourceSearchResponse{
|
||||
TotalHits: 100,
|
||||
},
|
||||
partialFailure: false,
|
||||
}
|
||||
assert.False(t, result.partialFailure)
|
||||
assert.NotNil(t, result.response)
|
||||
})
|
||||
|
||||
t.Run("partial failure result", func(t *testing.T) {
|
||||
result := &subIndexSearchResult{
|
||||
subIndexID: 3,
|
||||
err: assert.AnError,
|
||||
partialFailure: true,
|
||||
}
|
||||
assert.True(t, result.partialFailure)
|
||||
assert.Nil(t, result.response)
|
||||
assert.Error(t, result.err)
|
||||
})
|
||||
}
|
||||
|
||||
func TestSubIndexKey(t *testing.T) {
|
||||
key := SubIndexKey{
|
||||
NamespacedResource: NamespacedResource{
|
||||
Namespace: "org-1",
|
||||
Group: "dashboard.grafana.app",
|
||||
Resource: "dashboards",
|
||||
},
|
||||
SubIndexID: 42,
|
||||
}
|
||||
|
||||
t.Run("String representation", func(t *testing.T) {
|
||||
expected := "org-1/dashboard.grafana.app/dashboards/shard-42"
|
||||
assert.Equal(t, expected, key.String())
|
||||
})
|
||||
|
||||
t.Run("ToNamespacedResource", func(t *testing.T) {
|
||||
nsr := key.ToNamespacedResource()
|
||||
assert.Equal(t, "org-1", nsr.Namespace)
|
||||
assert.Equal(t, "dashboard.grafana.app", nsr.Group)
|
||||
assert.Equal(t, "dashboards", nsr.Resource)
|
||||
})
|
||||
}
|
||||
|
||||
func newTestDistributorServer() *distributorServer {
|
||||
return &distributorServer{
|
||||
tracing: noop.NewTracerProvider().Tracer("test"),
|
||||
log: log.New("test-distributor"),
|
||||
}
|
||||
}
|
||||
|
||||
func TestMergeResultsPagination(t *testing.T) {
|
||||
// Create mock results from 3 sub-indexes
|
||||
results := []*subIndexSearchResult{
|
||||
{
|
||||
subIndexID: 0,
|
||||
response: &resourcepb.ResourceSearchResponse{
|
||||
TotalHits: 10,
|
||||
Results: &resourcepb.ResourceTable{
|
||||
Columns: []*resourcepb.ResourceTableColumnDefinition{{Name: "title"}},
|
||||
Rows: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a1"}, Cells: [][]byte{[]byte("a1")}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "a2"}, Cells: [][]byte{[]byte("a2")}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
subIndexID: 1,
|
||||
response: &resourcepb.ResourceSearchResponse{
|
||||
TotalHits: 10,
|
||||
Results: &resourcepb.ResourceTable{
|
||||
Columns: []*resourcepb.ResourceTableColumnDefinition{{Name: "title"}},
|
||||
Rows: []*resourcepb.ResourceTableRow{
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b1"}, Cells: [][]byte{[]byte("b1")}},
|
||||
{Key: &resourcepb.ResourceKey{Namespace: "ns", Group: "g", Resource: "r", Name: "b2"}, Cells: [][]byte{[]byte("b2")}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
t.Run("pagination with limit", func(t *testing.T) {
|
||||
ds := newTestDistributorServer()
|
||||
req := &resourcepb.ResourceSearchRequest{
|
||||
Limit: 2,
|
||||
Offset: 0,
|
||||
Options: &resourcepb.ListOptions{
|
||||
Key: &resourcepb.ResourceKey{Namespace: "ns"},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := ds.mergeResults(context.Background(), results, req)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(20), resp.TotalHits) // 10 + 10
|
||||
assert.Len(t, resp.Results.Rows, 2)
|
||||
})
|
||||
|
||||
t.Run("pagination with offset", func(t *testing.T) {
|
||||
ds := newTestDistributorServer()
|
||||
req := &resourcepb.ResourceSearchRequest{
|
||||
Limit: 2,
|
||||
Offset: 2,
|
||||
Options: &resourcepb.ListOptions{
|
||||
Key: &resourcepb.ResourceKey{Namespace: "ns"},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := ds.mergeResults(context.Background(), results, req)
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, resp.Results.Rows, 2) // rows 3 and 4
|
||||
})
|
||||
}
|
||||
|
||||
func TestMergeResultsPartialFailure(t *testing.T) {
|
||||
results := []*subIndexSearchResult{
|
||||
{
|
||||
subIndexID: 0,
|
||||
response: &resourcepb.ResourceSearchResponse{
|
||||
TotalHits: 10,
|
||||
Results: &resourcepb.ResourceTable{
|
||||
Columns: []*resourcepb.ResourceTableColumnDefinition{{Name: "title"}},
|
||||
Rows: []*resourcepb.ResourceTableRow{},
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
subIndexID: 1,
|
||||
partialFailure: true,
|
||||
err: assert.AnError,
|
||||
},
|
||||
{
|
||||
subIndexID: 2,
|
||||
partialFailure: true,
|
||||
err: assert.AnError,
|
||||
},
|
||||
}
|
||||
|
||||
ds := newTestDistributorServer()
|
||||
req := &resourcepb.ResourceSearchRequest{
|
||||
Limit: 100,
|
||||
Options: &resourcepb.ListOptions{
|
||||
Key: &resourcepb.ResourceKey{Namespace: "ns"},
|
||||
},
|
||||
}
|
||||
|
||||
resp, err := ds.mergeResults(context.Background(), results, req)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Should have partial failure error
|
||||
require.NotNil(t, resp.Error)
|
||||
assert.Equal(t, int32(206), resp.Error.Code) // HTTP 206 Partial Content
|
||||
assert.Contains(t, resp.Error.Message, "2 of 3 sub-indexes failed")
|
||||
}
|
||||
Reference in New Issue
Block a user