@@ -7,33 +7,45 @@ import (
7
7
8
8
"github.com/google/uuid"
9
9
orderedmap "github.com/wk8/go-ordered-map/v2"
10
+ "go.opentelemetry.io/otel/attribute"
11
+ "go.opentelemetry.io/otel/codes"
12
+ "go.opentelemetry.io/otel/trace"
13
+ "go.opentelemetry.io/otel/trace/noop"
10
14
)
11
15
12
16
var ErrTooManyWaiters = fmt .Errorf ("rejecting request, too many waiters" )
13
17
14
18
type BoundedQueue struct {
15
- maxLimitBytes int64
19
+ maxLimitBytes int64
16
20
maxLimitWaiters int64
17
- currentBytes int64
18
- currentWaiters int64
19
- lock sync.Mutex
20
- waiters * orderedmap.OrderedMap [uuid.UUID , waiter ]
21
+ currentBytes int64
22
+ currentWaiters int64
23
+ lock sync.Mutex
24
+ waiters * orderedmap.OrderedMap [uuid.UUID , waiter ]
25
+ tracer trace.Tracer
21
26
}
22
27
23
28
type waiter struct {
24
- readyCh chan struct {}
29
+ readyCh chan struct {}
25
30
pendingBytes int64
26
- ID uuid.UUID
31
+ ID uuid.UUID
27
32
}
28
33
29
34
func NewBoundedQueue (maxLimitBytes , maxLimitWaiters int64 ) * BoundedQueue {
30
35
return & BoundedQueue {
31
- maxLimitBytes : maxLimitBytes ,
36
+ maxLimitBytes : maxLimitBytes ,
32
37
maxLimitWaiters : maxLimitWaiters ,
33
- waiters : orderedmap .New [uuid.UUID , waiter ](),
38
+ waiters : orderedmap .New [uuid.UUID , waiter ](),
39
+ tracer : noop .NewTracerProvider ().Tracer ("" ),
34
40
}
35
41
}
36
42
43
+ func NewTracedBoundedQueue (tp trace.TracerProvider , maxLimitBytes , maxLimitWaiters int64 ) * BoundedQueue {
44
+ bq := NewBoundedQueue (maxLimitBytes , maxLimitWaiters )
45
+ bq .tracer = tp .Tracer ("otel-arrow/admission" )
46
+ return bq
47
+ }
48
+
37
49
func (bq * BoundedQueue ) admit (pendingBytes int64 ) (bool , error ) {
38
50
bq .lock .Lock ()
39
51
defer bq .lock .Unlock ()
@@ -42,13 +54,13 @@ func (bq *BoundedQueue) admit(pendingBytes int64) (bool, error) {
42
54
return false , fmt .Errorf ("rejecting request, request size larger than configured limit" )
43
55
}
44
56
45
- if bq .currentBytes + pendingBytes <= bq .maxLimitBytes { // no need to wait to admit
57
+ if bq .currentBytes + pendingBytes <= bq .maxLimitBytes { // no need to wait to admit
46
58
bq .currentBytes += pendingBytes
47
59
return true , nil
48
60
}
49
61
50
62
// since we were unable to admit, check if we can wait.
51
- if bq .currentWaiters + 1 > bq .maxLimitWaiters { // too many waiters
63
+ if bq .currentWaiters + 1 > bq .maxLimitWaiters { // too many waiters
52
64
return false , ErrTooManyWaiters
53
65
}
54
66
@@ -66,7 +78,7 @@ func (bq *BoundedQueue) Acquire(ctx context.Context, pendingBytes int64) error {
66
78
// otherwise we need to wait for bytes to be released
67
79
curWaiter := waiter {
68
80
pendingBytes : pendingBytes ,
69
- readyCh : make (chan struct {}),
81
+ readyCh : make (chan struct {}),
70
82
}
71
83
72
84
bq .lock .Lock ()
@@ -84,6 +96,9 @@ func (bq *BoundedQueue) Acquire(ctx context.Context, pendingBytes int64) error {
84
96
}
85
97
86
98
bq .lock .Unlock ()
99
+ ctx , span := bq .tracer .Start (ctx , "admission_blocked" ,
100
+ trace .WithAttributes (attribute .Int64 ("pending" , pendingBytes )))
101
+ defer span .End ()
87
102
88
103
select {
89
104
case <- curWaiter .readyCh :
@@ -93,6 +108,7 @@ func (bq *BoundedQueue) Acquire(ctx context.Context, pendingBytes int64) error {
93
108
bq .lock .Lock ()
94
109
defer bq .lock .Unlock ()
95
110
err = fmt .Errorf ("context canceled: %w " , ctx .Err ())
111
+ span .SetStatus (codes .Error , "context canceled" )
96
112
97
113
_ , found := bq .waiters .Delete (curWaiter .ID )
98
114
if ! found {
@@ -121,7 +137,7 @@ func (bq *BoundedQueue) Release(pendingBytes int64) error {
121
137
next := bq .waiters .Oldest ()
122
138
nextWaiter := next .Value
123
139
nextKey := next .Key
124
- if bq .currentBytes + nextWaiter .pendingBytes <= bq .maxLimitBytes {
140
+ if bq .currentBytes + nextWaiter .pendingBytes <= bq .maxLimitBytes {
125
141
bq .currentBytes += nextWaiter .pendingBytes
126
142
bq .currentWaiters -= 1
127
143
close (nextWaiter .readyCh )
@@ -142,9 +158,9 @@ func (bq *BoundedQueue) Release(pendingBytes int64) error {
142
158
func (bq * BoundedQueue ) TryAcquire (pendingBytes int64 ) bool {
143
159
bq .lock .Lock ()
144
160
defer bq .lock .Unlock ()
145
- if bq .currentBytes + pendingBytes <= bq .maxLimitBytes {
161
+ if bq .currentBytes + pendingBytes <= bq .maxLimitBytes {
146
162
bq .currentBytes += pendingBytes
147
163
return true
148
164
}
149
165
return false
150
- }
166
+ }
0 commit comments