Sort traces in memory store to return most recent traces (#1394)

Fixes #1393

Signed-off-by: Jacob Marble <[email protected]>

plugin/storage/memory/memory.go

@@ -17,6 +17,7 @@ package memory
 import (
 	"context"
 	"errors"
+	"sort"
 	"sync"
 	"time"
 
@@ -188,13 +189,20 @@ func (m *Store) FindTraces(ctx context.Context, query *spanstore.TraceQueryParam
 	defer m.RUnlock()
 	var retMe []*model.Trace
 	for _, trace := range m.traces {
-		if len(retMe) >= query.NumTraces {
-			return retMe, nil
-		}
 		if m.validTrace(trace, query) {
 			retMe = append(retMe, trace)
 		}
 	}
+
+	// Query result order doesn't matter, as the query frontend will sort them anyway.
+	// However, if query.NumTraces < results, then we should return the newest traces.
+	if query.NumTraces > 0 && len(retMe) > query.NumTraces {
+		sort.Slice(retMe, func(i, j int) bool {
+			return retMe[i].Spans[0].StartTime.Before(retMe[j].Spans[0].StartTime)
+		})
+		retMe = retMe[len(retMe)-query.NumTraces:]
+	}
+
 	return retMe, nil
 }
 

plugin/storage/memory/memory_test.go

@@ -248,6 +248,55 @@ func TestStoreGetEmptyTraceSet(t *testing.T) {
 	})
 }
 
+func TestStoreFindTracesLimitGetsMostRecent(t *testing.T) {
+	storeSize, querySize := 100, 10
+
+	// This slice is in order from oldest to newest trace.
+	// Store keeps spans in a map, so storage order is effectively random.
+	// This ensures that query results include the most recent traces when limit < results.
+
+	var spans []*model.Span
+	for i := 0; i < storeSize; i++ {
+		spans = append(spans,
+			&model.Span{
+				TraceID:       model.NewTraceID(1, uint64(i)),
+				SpanID:        model.NewSpanID(1),
+				OperationName: "operationName",
+				Duration:      time.Second,
+				StartTime:     time.Unix(int64(i*24*60*60), 0),
+				Process: &model.Process{
+					ServiceName: "serviceName",
+				},
+			})
+	}
+
+	// Want the two most recent spans, not any two spans
+	var expectedTraces []*model.Trace
+	for _, span := range spans[storeSize-querySize:] {
+		trace := &model.Trace{
+			Spans: []*model.Span{span},
+		}
+		expectedTraces = append(expectedTraces, trace)
+	}
+
+	memStore := NewStore()
+	for _, span := range spans {
+		memStore.WriteSpan(span)
+	}
+
+	gotTraces, err := memStore.FindTraces(context.Background(), &spanstore.TraceQueryParameters{
+		ServiceName: "serviceName",
+		NumTraces:   querySize,
+	})
+
+	assert.NoError(t, err)
+	if assert.Len(t, gotTraces, len(expectedTraces)) {
+		for i := range gotTraces {
+			assert.EqualValues(t, expectedTraces[i].Spans[0].StartTime.Unix(), gotTraces[i].Spans[0].StartTime.Unix())
+		}
+	}
+}
+
 func TestStoreGetTrace(t *testing.T) {
 	testStruct := []struct {
 		query      *spanstore.TraceQueryParameters