forked from grafana/loki
-
Notifications
You must be signed in to change notification settings - Fork 0
/
parallel_chunk_fetch.go
82 lines (67 loc) · 1.73 KB
/
parallel_chunk_fetch.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
package util
import (
"context"
"sync"
otlog "github.com/opentracing/opentracing-go/log"
"github.com/frelon/loki/v2/pkg/util/spanlogger"
"github.com/frelon/loki/v2/pkg/storage/chunk"
)
var decodeContextPool = sync.Pool{
New: func() interface{} {
return chunk.NewDecodeContext()
},
}
// GetParallelChunks fetches chunks in parallel (up to maxParallel).
func GetParallelChunks(ctx context.Context, maxParallel int, chunks []chunk.Chunk, f func(context.Context, *chunk.DecodeContext, chunk.Chunk) (chunk.Chunk, error)) ([]chunk.Chunk, error) {
log, ctx := spanlogger.New(ctx, "GetParallelChunks")
defer log.Finish()
log.LogFields(otlog.Int("requested", len(chunks)))
if ctx.Err() != nil {
return nil, ctx.Err()
}
queuedChunks := make(chan chunk.Chunk)
go func() {
for _, c := range chunks {
queuedChunks <- c
}
close(queuedChunks)
}()
processedChunks := make(chan chunk.Chunk)
errors := make(chan error)
for i := 0; i < min(maxParallel, len(chunks)); i++ {
go func() {
decodeContext := decodeContextPool.Get().(*chunk.DecodeContext)
for c := range queuedChunks {
c, err := f(ctx, decodeContext, c)
if err != nil {
errors <- err
} else {
processedChunks <- c
}
}
decodeContextPool.Put(decodeContext)
}()
}
result := make([]chunk.Chunk, 0, len(chunks))
var lastErr error
for i := 0; i < len(chunks); i++ {
select {
case chunk := <-processedChunks:
result = append(result, chunk)
case err := <-errors:
lastErr = err
}
}
log.LogFields(otlog.Int("fetched", len(result)))
if lastErr != nil {
log.Error(lastErr)
}
// Return any chunks we did receive: a partial result may be useful
return result, lastErr
}
func min(a, b int) int {
if a < b {
return a
}
return b
}