github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/querier/queryrange/benchmarkutils_test.go (about)

     1  package queryrange
     2  
     3  import (
     4  	"sort"
     5  
     6  	"github.com/grafana/loki/pkg/logproto"
     7  )
     8  
     9  type entry struct {
    10  	entry  logproto.Entry
    11  	labels string
    12  }
    13  
    14  type byDirection struct {
    15  	direction logproto.Direction
    16  	entries   []entry
    17  }
    18  
    19  func (a byDirection) Len() int      { return len(a.entries) }
    20  func (a byDirection) Swap(i, j int) { a.entries[i], a.entries[j] = a.entries[j], a.entries[i] }
    21  func (a byDirection) Less(i, j int) bool {
    22  	e1, e2 := a.entries[i], a.entries[j]
    23  	if a.direction == logproto.BACKWARD {
    24  		switch {
    25  		case e1.entry.Timestamp.UnixNano() < e2.entry.Timestamp.UnixNano():
    26  			return false
    27  		case e1.entry.Timestamp.UnixNano() > e2.entry.Timestamp.UnixNano():
    28  			return true
    29  		default:
    30  			return e1.labels > e2.labels
    31  		}
    32  	}
    33  	switch {
    34  	case e1.entry.Timestamp.UnixNano() < e2.entry.Timestamp.UnixNano():
    35  		return true
    36  	case e1.entry.Timestamp.UnixNano() > e2.entry.Timestamp.UnixNano():
    37  		return false
    38  	default:
    39  		return e1.labels < e2.labels
    40  	}
    41  }
    42  
    43  func mergeStreams(resps []*LokiResponse, limit uint32, direction logproto.Direction) []logproto.Stream {
    44  	output := byDirection{
    45  		direction: direction,
    46  		entries:   []entry{},
    47  	}
    48  	for _, resp := range resps {
    49  		for _, stream := range resp.Data.Result {
    50  			for _, e := range stream.Entries {
    51  				output.entries = append(output.entries, entry{
    52  					entry:  e,
    53  					labels: stream.Labels,
    54  				})
    55  			}
    56  		}
    57  	}
    58  	sort.Sort(output)
    59  	// limit result
    60  	if len(output.entries) >= int(limit) {
    61  		output.entries = output.entries[:limit]
    62  	}
    63  
    64  	resultDict := map[string]*logproto.Stream{}
    65  	for _, e := range output.entries {
    66  		stream, ok := resultDict[e.labels]
    67  		if !ok {
    68  			stream = &logproto.Stream{
    69  				Labels:  e.labels,
    70  				Entries: []logproto.Entry{},
    71  			}
    72  			resultDict[e.labels] = stream
    73  		}
    74  		stream.Entries = append(stream.Entries, e.entry)
    75  
    76  	}
    77  	keys := make([]string, 0, len(resultDict))
    78  	for key := range resultDict {
    79  		keys = append(keys, key)
    80  	}
    81  	sort.Strings(keys)
    82  
    83  	result := make([]logproto.Stream, 0, len(resultDict))
    84  	for _, key := range keys {
    85  		result = append(result, *resultDict[key])
    86  	}
    87  
    88  	return result
    89  }