Skip to content

Commit f2bff77

Browse files
owen-dowen-d
authored
feat(block-scheduler): job tracking & offset commits (#15338)
1 parent 6316204 commit f2bff77

15 files changed

+767
-195
lines changed

‎pkg/blockbuilder/builder/builder.go

+8-8
Original file line numberDiff line numberDiff line change
@@ -250,13 +250,13 @@ func (i *BlockBuilder) runOne(ctx context.Context, workerID string) (bool, error
250250
logger := log.With(
251251
i.logger,
252252
"worker_id", workerID,
253-
"partition", job.Partition,
254-
"job_min_offset", job.Offsets.Min,
255-
"job_max_offset", job.Offsets.Max,
253+
"partition", job.Partition(),
254+
"job_min_offset", job.Offsets().Min,
255+
"job_max_offset", job.Offsets().Max,
256256
)
257257

258258
i.jobsMtx.Lock()
259-
i.inflightJobs[job.ID] = job
259+
i.inflightJobs[job.ID()] = job
260260
i.metrics.inflightJobs.Set(float64(len(i.inflightJobs)))
261261
i.jobsMtx.Unlock()
262262

@@ -284,7 +284,7 @@ func (i *BlockBuilder) runOne(ctx context.Context, workerID string) (bool, error
284284
}
285285

286286
i.jobsMtx.Lock()
287-
delete(i.inflightJobs, job.ID)
287+
delete(i.inflightJobs, job.ID())
288288
i.metrics.inflightJobs.Set(float64(len(i.inflightJobs)))
289289
i.jobsMtx.Unlock()
290290

@@ -315,15 +315,15 @@ func (i *BlockBuilder) processJob(ctx context.Context, job *types.Job, logger lo
315315
"load records",
316316
1,
317317
func(ctx context.Context) error {
318-
lastOffset, err = i.loadRecords(ctx, job.Partition, job.Offsets, inputCh)
318+
lastOffset, err = i.loadRecords(ctx, job.Partition(), job.Offsets(), inputCh)
319319
return err
320320
},
321321
func(ctx context.Context) error {
322322
level.Debug(logger).Log(
323323
"msg", "finished loading records",
324324
"ctx_error", ctx.Err(),
325325
"last_offset", lastOffset,
326-
"total_records", lastOffset-job.Offsets.Min,
326+
"total_records", lastOffset-job.Offsets().Min,
327327
)
328328
close(inputCh)
329329
return nil
@@ -488,7 +488,7 @@ func (i *BlockBuilder) processJob(ctx context.Context, job *types.Job, logger lo
488488
}
489489
}
490490

491-
if lastOffset <= job.Offsets.Min {
491+
if lastOffset <= job.Offsets().Min {
492492
return lastOffset, nil
493493
}
494494

‎pkg/blockbuilder/scheduler/prioritiy_queue_test.go

+127-17
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77
)
88

99
func TestPriorityQueue(t *testing.T) {
10-
t.Run("operations", func(t *testing.T) {
10+
t.Run("basic operations", func(t *testing.T) {
1111
tests := []struct {
1212
name string
1313
input []int
@@ -33,16 +33,14 @@ func TestPriorityQueue(t *testing.T) {
3333
input: []int{3, 1, 2},
3434
wantPops: []int{1, 2, 3},
3535
},
36-
{
37-
name: "duplicate elements",
38-
input: []int{2, 1, 2, 1},
39-
wantPops: []int{1, 1, 2, 2},
40-
},
4136
}
4237

4338
for _, tt := range tests {
4439
t.Run(tt.name, func(t *testing.T) {
45-
pq := NewPriorityQueue[int](func(a, b int) bool { return a < b })
40+
pq := NewPriorityQueue[int, int](
41+
func(a, b int) bool { return a < b },
42+
func(v int) int { return v },
43+
)
4644
require.Equal(t, 0, pq.Len())
4745

4846
// Push all elements
@@ -69,15 +67,73 @@ func TestPriorityQueue(t *testing.T) {
6967
}
7068
})
7169

70+
t.Run("key operations", func(t *testing.T) {
71+
type Job struct {
72+
ID string
73+
Priority int
74+
}
75+
76+
pq := NewPriorityQueue[string, Job](
77+
func(a, b Job) bool { return a.Priority < b.Priority },
78+
func(j Job) string { return j.ID },
79+
)
80+
81+
// Test Push with duplicate key
82+
job1 := Job{ID: "job1", Priority: 1}
83+
job1Updated := Job{ID: "job1", Priority: 3}
84+
job2 := Job{ID: "job2", Priority: 2}
85+
86+
pq.Push(job1)
87+
require.Equal(t, 1, pq.Len())
88+
89+
// Push with same key should update
90+
pq.Push(job1Updated)
91+
require.Equal(t, 1, pq.Len())
92+
93+
// Verify updated priority
94+
v, ok := pq.Lookup("job1")
95+
require.True(t, ok)
96+
require.Equal(t, job1Updated, v)
97+
98+
// Test Remove
99+
pq.Push(job2)
100+
v, ok = pq.Remove("job1")
101+
require.True(t, ok)
102+
require.Equal(t, job1Updated, v)
103+
require.Equal(t, 1, pq.Len())
104+
105+
// Test UpdatePriority
106+
newJob2 := Job{ID: "job2", Priority: 4}
107+
ok = pq.UpdatePriority("job2", newJob2)
108+
require.True(t, ok)
109+
110+
v, ok = pq.Lookup("job2")
111+
require.True(t, ok)
112+
require.Equal(t, newJob2, v)
113+
114+
// Test non-existent key operations
115+
v, ok = pq.Lookup("nonexistent")
116+
require.False(t, ok)
117+
require.Zero(t, v)
118+
119+
v, ok = pq.Remove("nonexistent")
120+
require.False(t, ok)
121+
require.Zero(t, v)
122+
123+
ok = pq.UpdatePriority("nonexistent", Job{})
124+
require.False(t, ok)
125+
})
126+
72127
t.Run("custom type", func(t *testing.T) {
73128
type Job struct {
74129
ID string
75130
Priority int
76131
}
77132

78-
pq := NewPriorityQueue[Job](func(a, b Job) bool {
79-
return a.Priority < b.Priority
80-
})
133+
pq := NewPriorityQueue[string, Job](
134+
func(a, b Job) bool { return a.Priority < b.Priority },
135+
func(j Job) string { return j.ID },
136+
)
81137

82138
jobs := []Job{
83139
{ID: "high", Priority: 3},
@@ -102,25 +158,28 @@ func TestPriorityQueue(t *testing.T) {
102158
})
103159

104160
t.Run("mixed operations", func(t *testing.T) {
105-
pq := NewPriorityQueue[int](func(a, b int) bool { return a < b })
161+
pq := NewPriorityQueue[int, int](
162+
func(a, b int) bool { return a < b },
163+
func(v int) int { return v },
164+
)
106165

107166
// Push some elements
108167
pq.Push(3)
109168
pq.Push(1)
110-
require.Equal(t, 2, pq.Len())
169+
pq.Push(4)
111170

112-
// Pop lowest
171+
// Pop an element
113172
v, ok := pq.Pop()
114173
require.True(t, ok)
115174
require.Equal(t, 1, v)
116175

117176
// Push more elements
118177
pq.Push(2)
119-
pq.Push(4)
178+
pq.Push(5)
120179

121-
// Verify remaining elements come out in order
122-
want := []int{2, 3, 4}
123-
got := make([]int, 0, 3)
180+
// Pop remaining elements and verify order
181+
want := []int{2, 3, 4, 5}
182+
got := make([]int, 0, len(want))
124183
for range want {
125184
v, ok := pq.Pop()
126185
require.True(t, ok)
@@ -191,3 +250,54 @@ func TestCircularBuffer(t *testing.T) {
191250
})
192251
}
193252
}
253+
254+
func TestCircularBufferLookup(t *testing.T) {
255+
t.Run("empty buffer", func(t *testing.T) {
256+
cb := NewCircularBuffer[int](5)
257+
_, ok := cb.Lookup(func(i int) bool { return i == 1 })
258+
require.False(t, ok)
259+
})
260+
261+
t.Run("single element", func(t *testing.T) {
262+
cb := NewCircularBuffer[int](5)
263+
cb.Push(1)
264+
v, ok := cb.Lookup(func(i int) bool { return i == 1 })
265+
require.True(t, ok)
266+
require.Equal(t, 1, v)
267+
})
268+
269+
t.Run("multiple elements", func(t *testing.T) {
270+
cb := NewCircularBuffer[int](5)
271+
for i := 1; i <= 3; i++ {
272+
cb.Push(i)
273+
}
274+
v, ok := cb.Lookup(func(i int) bool { return i == 2 })
275+
require.True(t, ok)
276+
require.Equal(t, 2, v)
277+
})
278+
279+
t.Run("wrapped buffer", func(t *testing.T) {
280+
cb := NewCircularBuffer[int](3)
281+
// Push 5 elements into a buffer of size 3, causing wrap-around
282+
for i := 1; i <= 5; i++ {
283+
cb.Push(i)
284+
}
285+
// Buffer should now contain [4,5,3] with head at index 2
286+
v, ok := cb.Lookup(func(i int) bool { return i == 4 })
287+
require.True(t, ok)
288+
require.Equal(t, 4, v)
289+
290+
// Element that was evicted should not be found
291+
_, ok = cb.Lookup(func(i int) bool { return i == 1 })
292+
require.False(t, ok)
293+
})
294+
295+
t.Run("no match", func(t *testing.T) {
296+
cb := NewCircularBuffer[int](5)
297+
for i := 1; i <= 3; i++ {
298+
cb.Push(i)
299+
}
300+
_, ok := cb.Lookup(func(i int) bool { return i == 99 })
301+
require.False(t, ok)
302+
})
303+
}

0 commit comments

Comments
 (0)