this repo has no description
integrate consumer pool into stream handler (#179)
Add an easy way to process the event stream in parallel
authored by
Whyrusleeping
and committed by
GitHub
4a875abc
7f792c32
+159
-100
+5
-2
bgs/fedmgr.go
+5
-2
bgs/fedmgr.go
···
231
231
ctx, cancel := context.WithCancel(ctx)
232
232
defer cancel()
233
233
234
-
return events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
234
+
rsc := &events.RepoStreamCallbacks{
235
235
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
236
236
log.Infow("got remote repo event", "host", host.Host, "repo", evt.Repo, "seq", evt.Seq)
237
237
if err := s.cb(context.TODO(), host, &events.XRPCStreamEvent{
···
300
300
Error: func(errf *events.ErrorFrame) error {
301
301
return fmt.Errorf("error frame: %s: %s", errf.Error, errf.Message)
302
302
},
303
-
})
303
+
}
304
+
305
+
// TODO: probably make this use the parallel handler after some thought
306
+
return events.HandleRepoStream(ctx, con, &events.SequentialScheduler{rsc.EventHandler})
304
307
}
305
308
306
309
func (s *Slurper) updateCursor(host *models.PDS, curs int64) error {
+6
-4
cmd/gosky/debug.go
+6
-4
cmd/gosky/debug.go
···
69
69
var match *comatproto.SyncSubscribeRepos_Commit
70
70
71
71
ctx := context.TODO()
72
-
err = events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
72
+
rsc := &events.RepoStreamCallbacks{
73
73
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
74
74
n := int64(n)
75
75
if evt.Seq == n {
···
89
89
Error: func(evt *events.ErrorFrame) error {
90
90
return fmt.Errorf("%s: %s", evt.Error, evt.Message)
91
91
},
92
-
})
92
+
}
93
93
94
+
err = events.HandleRepoStream(ctx, con, &events.SequentialScheduler{rsc.EventHandler})
94
95
if err != errFoundIt {
95
96
return err
96
97
}
···
190
191
infos := make(map[string]*eventInfo)
191
192
192
193
ctx := context.TODO()
193
-
err = events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
194
+
rsc := &events.RepoStreamCallbacks{
194
195
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
195
196
196
197
fmt.Printf("\rChecking seq: %d ", evt.Seq)
···
243
244
Error: func(evt *events.ErrorFrame) error {
244
245
return fmt.Errorf("%s: %s", evt.Error, evt.Message)
245
246
},
246
-
})
247
+
}
248
+
err = events.HandleRepoStream(ctx, con, &events.SequentialScheduler{rsc.EventHandler})
247
249
if err != nil {
248
250
return err
249
251
}
+3
-2
cmd/gosky/main.go
+3
-2
cmd/gosky/main.go
···
938
938
_ = con.Close()
939
939
}()
940
940
941
-
return events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
941
+
rsc := &events.RepoStreamCallbacks{
942
942
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
943
943
if jsonfmt {
944
944
b, err := json.Marshal(evt)
···
984
984
Error: func(errf *events.ErrorFrame) error {
985
985
return fmt.Errorf("error frame: %s: %s", errf.Error, errf.Message)
986
986
},
987
-
})
987
+
}
988
+
return events.HandleRepoStream(ctx, con, &events.SequentialScheduler{rsc.EventHandler})
988
989
},
989
990
}
990
991
+6
-4
cmd/gosky/streamdiff.go
+6
-4
cmd/gosky/streamdiff.go
···
40
40
defer cancel()
41
41
42
42
go func() {
43
-
err = events.HandleRepoStream(ctx, cona, &events.RepoStreamCallbacks{
43
+
rsc := &events.RepoStreamCallbacks{
44
44
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
45
45
sd.PushA(&events.XRPCStreamEvent{
46
46
RepoCommit: evt,
···
54
54
Error: func(evt *events.ErrorFrame) error {
55
55
return fmt.Errorf("%s: %s", evt.Error, evt.Message)
56
56
},
57
-
})
57
+
}
58
+
err = events.HandleRepoStream(ctx, cona, &events.SequentialScheduler{rsc.EventHandler})
58
59
if err != nil {
59
60
log.Errorf("stream A failed: %s", err)
60
61
}
61
62
}()
62
63
63
64
go func() {
64
-
err = events.HandleRepoStream(ctx, conb, &events.RepoStreamCallbacks{
65
+
rsc := &events.RepoStreamCallbacks{
65
66
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
66
67
sd.PushB(&events.XRPCStreamEvent{
67
68
RepoCommit: evt,
···
75
76
Error: func(evt *events.ErrorFrame) error {
76
77
return fmt.Errorf("%s: %s", evt.Error, evt.Message)
77
78
},
78
-
})
79
+
}
80
+
err = events.HandleRepoStream(ctx, conb, &events.SequentialScheduler{rsc.EventHandler})
79
81
if err != nil {
80
82
log.Errorf("stream A failed: %s", err)
81
83
}
+68
-33
events/consumer.go
+68
-33
events/consumer.go
···
22
22
Error func(evt *ErrorFrame) error
23
23
}
24
24
25
-
func HandleRepoStream(ctx context.Context, con *websocket.Conn, cbs *RepoStreamCallbacks) error {
25
+
func (rsc *RepoStreamCallbacks) EventHandler(ctx context.Context, xev *XRPCStreamEvent) error {
26
+
switch {
27
+
case xev.RepoCommit != nil && rsc.RepoCommit != nil:
28
+
return rsc.RepoCommit(xev.RepoCommit)
29
+
case xev.RepoHandle != nil && rsc.RepoHandle != nil:
30
+
return rsc.RepoHandle(xev.RepoHandle)
31
+
case xev.RepoInfo != nil && rsc.RepoInfo != nil:
32
+
return rsc.RepoInfo(xev.RepoInfo)
33
+
case xev.RepoMigrate != nil && rsc.RepoMigrate != nil:
34
+
return rsc.RepoMigrate(xev.RepoMigrate)
35
+
case xev.RepoTombstone != nil && rsc.RepoTombstone != nil:
36
+
return rsc.RepoTombstone(xev.RepoTombstone)
37
+
case xev.LabelLabels != nil && rsc.LabelLabels != nil:
38
+
return rsc.LabelLabels(xev.LabelLabels)
39
+
case xev.LabelInfo != nil && rsc.LabelInfo != nil:
40
+
return rsc.LabelInfo(xev.LabelInfo)
41
+
case xev.Error != nil && rsc.Error != nil:
42
+
return rsc.Error(xev.Error)
43
+
default:
44
+
return fmt.Errorf("no know event in XRPCStreamEvent object")
45
+
}
46
+
}
47
+
48
+
func HandleRepoStream(ctx context.Context, con *websocket.Conn, sched Scheduler) error {
26
49
ctx, cancel := context.WithCancel(ctx)
27
50
defer cancel()
28
51
···
76
99
77
100
lastSeq = evt.Seq
78
101
79
-
if cbs.RepoCommit != nil {
80
-
if err := cbs.RepoCommit(&evt); err != nil {
81
-
return err
82
-
}
83
-
} else {
84
-
log.Warnf("received repo commit event with nil commit object (seq %d)", evt.Seq)
102
+
if err := sched.AddWork(ctx, evt.Repo, &XRPCStreamEvent{
103
+
RepoCommit: &evt,
104
+
}); err != nil {
105
+
return err
85
106
}
86
107
case "#handle":
87
108
var evt comatproto.SyncSubscribeRepos_Handle
···
89
110
return err
90
111
}
91
112
92
-
if cbs.RepoHandle != nil {
93
-
if err := cbs.RepoHandle(&evt); err != nil {
94
-
return err
95
-
}
113
+
if evt.Seq < lastSeq {
114
+
log.Errorf("Got events out of order from stream (seq = %d, prev = %d)", evt.Seq, lastSeq)
115
+
}
116
+
lastSeq = evt.Seq
117
+
118
+
if err := sched.AddWork(ctx, evt.Did, &XRPCStreamEvent{
119
+
RepoHandle: &evt,
120
+
}); err != nil {
121
+
return err
96
122
}
97
123
case "#info":
98
124
// TODO: this might also be a LabelInfo (as opposed to RepoInfo)
···
101
127
return err
102
128
}
103
129
104
-
if cbs.RepoInfo != nil {
105
-
if err := cbs.RepoInfo(&evt); err != nil {
106
-
return err
107
-
}
130
+
if err := sched.AddWork(ctx, "", &XRPCStreamEvent{
131
+
RepoInfo: &evt,
132
+
}); err != nil {
133
+
return err
108
134
}
109
135
case "#migrate":
110
136
var evt comatproto.SyncSubscribeRepos_Migrate
···
112
138
return err
113
139
}
114
140
115
-
if cbs.RepoMigrate != nil {
116
-
if err := cbs.RepoMigrate(&evt); err != nil {
117
-
return err
118
-
}
141
+
if evt.Seq < lastSeq {
142
+
log.Errorf("Got events out of order from stream (seq = %d, prev = %d)", evt.Seq, lastSeq)
143
+
}
144
+
lastSeq = evt.Seq
145
+
146
+
if err := sched.AddWork(ctx, evt.Did, &XRPCStreamEvent{
147
+
RepoMigrate: &evt,
148
+
}); err != nil {
149
+
return err
119
150
}
120
151
case "#tombstone":
121
152
var evt comatproto.SyncSubscribeRepos_Tombstone
···
123
154
return err
124
155
}
125
156
126
-
if cbs.RepoMigrate != nil {
127
-
if err := cbs.RepoTombstone(&evt); err != nil {
128
-
return err
129
-
}
157
+
if evt.Seq < lastSeq {
158
+
log.Errorf("Got events out of order from stream (seq = %d, prev = %d)", evt.Seq, lastSeq)
159
+
}
160
+
lastSeq = evt.Seq
161
+
162
+
if err := sched.AddWork(ctx, evt.Did, &XRPCStreamEvent{
163
+
RepoTombstone: &evt,
164
+
}); err != nil {
165
+
return err
130
166
}
131
167
case "#labebatch":
132
168
var evt label.SubscribeLabels_Labels
···
140
176
141
177
lastSeq = evt.Seq
142
178
143
-
if cbs.LabelLabels != nil {
144
-
if err := cbs.LabelLabels(&evt); err != nil {
145
-
return err
146
-
}
147
-
} else {
148
-
log.Warnf("received label event with nil append object (seq %d)", evt.Seq)
179
+
if err := sched.AddWork(ctx, "", &XRPCStreamEvent{
180
+
LabelLabels: &evt,
181
+
}); err != nil {
182
+
return err
149
183
}
150
184
}
151
185
···
155
189
return err
156
190
}
157
191
158
-
if cbs.Error != nil {
159
-
if err := cbs.Error(&errframe); err != nil {
160
-
return err
161
-
}
192
+
if err := sched.AddWork(ctx, "", &XRPCStreamEvent{
193
+
Error: &errframe,
194
+
}); err != nil {
195
+
return err
162
196
}
197
+
163
198
default:
164
199
return fmt.Errorf("unrecognized event stream type: %d", header.Op)
165
200
}
+21
-7
events/parallel.go
+21
-7
events/parallel.go
···
5
5
"sync"
6
6
)
7
7
8
-
type ConsumerPool struct {
8
+
type Scheduler interface {
9
+
AddWork(ctx context.Context, repo string, val *XRPCStreamEvent) error
10
+
}
11
+
12
+
type SequentialScheduler struct {
13
+
Do func(context.Context, *XRPCStreamEvent) error
14
+
}
15
+
16
+
func (s *SequentialScheduler) AddWork(ctx context.Context, repo string, val *XRPCStreamEvent) error {
17
+
return s.Do(ctx, val)
18
+
}
19
+
20
+
type ParallelConsumerPool struct {
9
21
maxConcurrency int
10
22
maxQueue int
11
23
12
-
do func(*XRPCStreamEvent)
24
+
do func(context.Context, *XRPCStreamEvent) error
13
25
14
26
feeder chan *consumerTask
15
27
···
17
29
active map[string][]*consumerTask
18
30
}
19
31
20
-
func NewConsumerPool(maxC, maxQ int, do func(*XRPCStreamEvent)) *ConsumerPool {
21
-
p := &ConsumerPool{
32
+
func NewConsumerPool(maxC, maxQ int, do func(context.Context, *XRPCStreamEvent) error) *ParallelConsumerPool {
33
+
p := &ParallelConsumerPool{
22
34
maxConcurrency: maxC,
23
35
maxQueue: maxQ,
24
36
···
40
52
val *XRPCStreamEvent
41
53
}
42
54
43
-
func (p *ConsumerPool) Add(ctx context.Context, repo string, val *XRPCStreamEvent) error {
55
+
func (p *ParallelConsumerPool) AddWork(ctx context.Context, repo string, val *XRPCStreamEvent) error {
44
56
t := &consumerTask{
45
57
repo: repo,
46
58
val: val,
···
65
77
}
66
78
}
67
79
68
-
func (p *ConsumerPool) worker() {
80
+
func (p *ParallelConsumerPool) worker() {
69
81
for work := range p.feeder {
70
82
for work != nil {
71
-
p.do(work.val)
83
+
if err := p.do(context.TODO(), work.val); err != nil {
84
+
log.Errorf("event handler failed: %s", err)
85
+
}
72
86
73
87
p.lk.Lock()
74
88
rem, ok := p.active[work.repo]
+40
-42
events/repostream.go
+40
-42
events/repostream.go
···
1
1
package events
2
2
3
3
import (
4
-
"bytes"
5
4
"context"
6
-
"fmt"
7
5
8
-
comatproto "github.com/bluesky-social/indigo/api/atproto"
9
-
lexutil "github.com/bluesky-social/indigo/lex/util"
10
-
"github.com/bluesky-social/indigo/repo"
11
6
"github.com/bluesky-social/indigo/repomgr"
12
7
13
8
"github.com/gorilla/websocket"
···
16
11
17
12
type LiteStreamHandleFunc func(op repomgr.EventKind, seq int64, path string, did string, rcid *cid.Cid, rec any) error
18
13
19
-
func ConsumeRepoStreamLite(ctx context.Context, con *websocket.Conn, cb LiteStreamHandleFunc) error {
20
-
return HandleRepoStream(ctx, con, &RepoStreamCallbacks{
21
-
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
22
-
if evt.TooBig {
23
-
log.Errorf("skipping too big events for now: %d", evt.Seq)
24
-
return nil
25
-
}
26
-
r, err := repo.ReadRepoFromCar(ctx, bytes.NewReader(evt.Blocks))
27
-
if err != nil {
28
-
return fmt.Errorf("reading repo from car (seq: %d, len: %d): %w", evt.Seq, len(evt.Blocks), err)
29
-
}
14
+
func ConsumeRepoStreamLite2(ctx context.Context, con *websocket.Conn, cb LiteStreamHandleFunc) error {
15
+
/*
16
+
return HandleRepoStream(ctx, con, &RepoStreamCallbacks{
17
+
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
18
+
if evt.TooBig {
19
+
log.Errorf("skipping too big events for now: %d", evt.Seq)
20
+
return nil
21
+
}
22
+
r, err := repo.ReadRepoFromCar(ctx, bytes.NewReader(evt.Blocks))
23
+
if err != nil {
24
+
return fmt.Errorf("reading repo from car (seq: %d, len: %d): %w", evt.Seq, len(evt.Blocks), err)
25
+
}
30
26
31
-
for _, op := range evt.Ops {
32
-
ek := repomgr.EventKind(op.Action)
33
-
switch ek {
34
-
case repomgr.EvtKindCreateRecord, repomgr.EvtKindUpdateRecord:
35
-
rc, rec, err := r.GetRecord(ctx, op.Path)
36
-
if err != nil {
37
-
e := fmt.Errorf("getting record %s (%s) within seq %d for %s: %w", op.Path, *op.Cid, evt.Seq, evt.Repo, err)
38
-
log.Error(e)
39
-
continue
40
-
}
27
+
for _, op := range evt.Ops {
28
+
ek := repomgr.EventKind(op.Action)
29
+
switch ek {
30
+
case repomgr.EvtKindCreateRecord, repomgr.EvtKindUpdateRecord:
31
+
rc, rec, err := r.GetRecord(ctx, op.Path)
32
+
if err != nil {
33
+
e := fmt.Errorf("getting record %s (%s) within seq %d for %s: %w", op.Path, *op.Cid, evt.Seq, evt.Repo, err)
34
+
log.Error(e)
35
+
continue
36
+
}
41
37
42
-
if lexutil.LexLink(rc) != *op.Cid {
43
-
// TODO: do we even error here?
44
-
return fmt.Errorf("mismatch in record and op cid: %s != %s", rc, *op.Cid)
45
-
}
38
+
if lexutil.LexLink(rc) != *op.Cid {
39
+
// TODO: do we even error here?
40
+
return fmt.Errorf("mismatch in record and op cid: %s != %s", rc, *op.Cid)
41
+
}
46
42
47
-
if err := cb(ek, evt.Seq, op.Path, evt.Repo, &rc, rec); err != nil {
48
-
log.Errorf("event consumer callback (%s): %s", ek, err)
49
-
continue
50
-
}
43
+
if err := cb(ek, evt.Seq, op.Path, evt.Repo, &rc, rec); err != nil {
44
+
log.Errorf("event consumer callback (%s): %s", ek, err)
45
+
continue
46
+
}
51
47
52
-
case repomgr.EvtKindDeleteRecord:
53
-
if err := cb(ek, evt.Seq, op.Path, evt.Repo, nil, nil); err != nil {
54
-
log.Errorf("event consumer callback (%s): %s", ek, err)
55
-
continue
48
+
case repomgr.EvtKindDeleteRecord:
49
+
if err := cb(ek, evt.Seq, op.Path, evt.Repo, nil, nil); err != nil {
50
+
log.Errorf("event consumer callback (%s): %s", ek, err)
51
+
continue
52
+
}
56
53
}
57
54
}
58
-
}
59
-
return nil
60
-
},
61
-
})
55
+
return nil
56
+
},
57
+
})
58
+
*/
59
+
return nil
62
60
}
+4
-2
search/server.go
+4
-2
search/server.go
···
133
133
return fmt.Errorf("events dial failed: %w", err)
134
134
}
135
135
136
-
return events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
136
+
rsc := &events.RepoStreamCallbacks{
137
137
RepoCommit: func(evt *comatproto.SyncSubscribeRepos_Commit) error {
138
138
if evt.TooBig && evt.Prev != nil {
139
139
log.Errorf("skipping non-genesis too big events for now: %d", evt.Seq)
···
193
193
}
194
194
return nil
195
195
},
196
-
})
196
+
}
197
+
198
+
return events.HandleRepoStream(ctx, con, events.NewConsumerPool(8, 32, rsc.EventHandler))
197
199
}
198
200
199
201
func (s *Server) handleOp(ctx context.Context, op repomgr.EventKind, seq int64, path string, did string, rcid *cid.Cid, rec any) error {
+3
-2
testing/labelmaker_fakedata_test.go
+3
-2
testing/labelmaker_fakedata_test.go
···
100
100
}()
101
101
102
102
go func() {
103
-
if err := events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
103
+
rsc := &events.RepoStreamCallbacks{
104
104
LabelLabels: func(evt *label.SubscribeLabels_Labels) error {
105
105
fmt.Println("received event: ", evt.Seq)
106
106
es.lk.Lock()
···
108
108
es.lk.Unlock()
109
109
return nil
110
110
},
111
-
}); err != nil {
111
+
}
112
+
if err := events.HandleRepoStream(ctx, con, &events.SequentialScheduler{rsc.EventHandler}); err != nil {
112
113
fmt.Println(err)
113
114
}
114
115
}()
+3
-2
testing/utils.go
+3
-2
testing/utils.go
···
467
467
}()
468
468
469
469
go func() {
470
-
if err := events.HandleRepoStream(ctx, con, &events.RepoStreamCallbacks{
470
+
rsc := &events.RepoStreamCallbacks{
471
471
RepoCommit: func(evt *atproto.SyncSubscribeRepos_Commit) error {
472
472
fmt.Println("received event: ", evt.Seq, evt.Repo)
473
473
es.lk.Lock()
···
482
482
es.lk.Unlock()
483
483
return nil
484
484
},
485
-
}); err != nil {
485
+
}
486
+
if err := events.HandleRepoStream(ctx, con, &events.SequentialScheduler{rsc.EventHandler}); err != nil {
486
487
fmt.Println(err)
487
488
}
488
489
}()