forked from stellar/go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
stream_handler.go
88 lines (76 loc) · 2.05 KB
/
stream_handler.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
83
84
85
86
87
88
package sse
import (
"net/http"
"github.com/aiblocks/go/services/millennium/internal/ledger"
"github.com/aiblocks/go/support/errors"
"github.com/stellar/throttled"
)
type LedgerSourceFactory interface {
Get() ledger.Source
}
// StreamHandler represents a stream handling action
type StreamHandler struct {
RateLimiter *throttled.HTTPRateLimiter
LedgerSourceFactory LedgerSourceFactory
}
// GenerateEventsFunc generates a slice of sse.Event which are sent via
// streaming.
type GenerateEventsFunc func() ([]Event, error)
// ServeStream handles a SSE requests, sending data every time there is a new
// ledger.
func (handler StreamHandler) ServeStream(
w http.ResponseWriter,
r *http.Request,
limit int,
generateEvents GenerateEventsFunc,
) {
ctx := r.Context()
stream := NewStream(ctx, w)
stream.SetLimit(limit)
ledgerSource := handler.LedgerSourceFactory.Get()
defer ledgerSource.Close()
currentLedgerSequence := ledgerSource.CurrentLedger()
for {
// Rate limit the request if it's a call to stream since it queries the DB every second. See
// https://github.com/aiblocks/go/issues/715 for more details.
rateLimiter := handler.RateLimiter
if rateLimiter != nil {
limited, _, err := rateLimiter.RateLimiter.RateLimit(rateLimiter.VaryBy.Key(r), 1)
if err != nil {
stream.Err(errors.Wrap(err, "RateLimiter error"))
return
}
if limited {
stream.Err(ErrRateLimited)
return
}
}
events, err := generateEvents()
if err != nil {
stream.Err(err)
return
}
for _, event := range events {
if limit <= 0 {
break
}
stream.Send(event)
limit--
}
if limit <= 0 {
stream.Done()
return
}
// Manually send the preamble in case there are no data events in SSE to trigger a stream.Send call.
// This method is called every iteration of the loop, but is protected by a sync.Once variable so it's
// only executed once.
stream.Init()
select {
case currentLedgerSequence = <-ledgerSource.NextLedger(currentLedgerSequence):
continue
case <-ctx.Done():
stream.Done()
return
}
}
}