@@ -4,97 +4,107 @@ import (
44 "context"
55 "time"
66
7+ "github.com/mr-tron/base58"
78 "github.com/pkg/errors"
8-
9- "github.com/code-payments/code-server/pkg/solana/cvm"
109 "google.golang.org/grpc"
1110 "google.golang.org/grpc/credentials/insecure"
11+ "google.golang.org/grpc/metadata"
1212
1313 geyserpb "github.com/code-payments/code-vm-indexer/generated/geyser/v1"
14+
15+ "github.com/code-payments/code-server/pkg/solana/cvm"
1416)
1517
1618const (
1719 defaultStreamSubscriptionTimeout = time .Minute
1820)
1921
2022var (
21- ErrSubscriptionFallenBehind = errors .New ("subscription stream fell behind" )
22- ErrTimeoutReceivingUpdate = errors .New ("timed out receiving update" )
23+ ErrTimeoutReceivingUpdate = errors .New ("timed out receiving update" )
2324)
2425
25- func newGeyserClient (ctx context.Context , endpoint string ) (geyserpb.GeyserClient , error ) {
26- conn , err := grpc .Dial (endpoint , grpc .WithTransportCredentials (insecure .NewCredentials ()))
26+ func newGeyserClient (endpoint , xToken string ) (geyserpb.GeyserClient , error ) {
27+ opts := []grpc.DialOption {grpc .WithTransportCredentials (insecure .NewCredentials ())}
28+ if len (xToken ) > 0 {
29+ opts = append (
30+ opts ,
31+ grpc .WithUnaryInterceptor (newXTokenUnaryClientInterceptor (xToken )),
32+ grpc .WithStreamInterceptor (newXTokenStreamClientInterceptor (xToken )),
33+ )
34+ }
35+
36+ conn , err := grpc .NewClient (endpoint , opts ... )
2737 if err != nil {
2838 return nil , err
2939 }
3040
3141 client := geyserpb .NewGeyserClient (conn )
3242
33- // Unfortunately the RPCs we use no longer support hearbeats. We'll let each
34- // individual subscriber determine what an appropriate timeout to receive a
35- // message should be.
36- /*
37- heartbeatResp, err := client.GetHeartbeatInterval(ctx, &geyserpb.EmptyRequest{})
38- if err != nil {
39- return nil, 0, errors.Wrap(err, "error getting heartbeat interval")
40- }
41-
42- heartbeatTimeout := time.Duration(2 * heartbeatResp.HeartbeatIntervalMs * uint64(time.Millisecond))
43- */
44-
4543 return client , nil
4644}
4745
48- func boundedProgramUpdateRecv (ctx context.Context , streamer geyserpb.Geyser_SubscribeProgramUpdatesClient , timeout time.Duration ) (update * geyserpb.TimestampedAccountUpdate , err error ) {
46+ func boundedRecv (ctx context.Context , streamer geyserpb.Geyser_SubscribeClient , timeout time.Duration ) (update * geyserpb.SubscribeUpdate , err error ) {
4947 done := make (chan struct {})
5048 go func () {
5149 update , err = streamer .Recv ()
5250 close (done )
5351 }()
5452
5553 select {
54+ case <- ctx .Done ():
55+ return nil , ctx .Err ()
5656 case <- time .After (timeout ):
5757 return nil , ErrTimeoutReceivingUpdate
5858 case <- done :
5959 return update , err
6060 }
6161}
6262
63- func (w * Worker ) subscribeToProgramUpdatesFromGeyser (ctx context.Context , endpoint string ) error {
63+ func (w * Worker ) subscribeToProgramUpdatesFromGeyser (ctx context.Context , endpoint , xToken string ) error {
6464 log := w .log .WithField ("method" , "subscribeToProgramUpdatesFromGeyser" )
6565 log .Debug ("subscription started" )
6666
6767 defer func () {
6868 log .Debug ("subscription stopped" )
6969 }()
7070
71- client , err := newGeyserClient (ctx , endpoint )
71+ client , err := newGeyserClient (endpoint , xToken )
7272 if err != nil {
7373 return errors .Wrap (err , "error creating client" )
7474 }
7575
76- streamer , err := client .SubscribeProgramUpdates (ctx , & geyserpb.SubscribeProgramsUpdatesRequest {
77- Programs : [][]byte {cvm .PROGRAM_ID },
78- })
76+ streamer , err := client .Subscribe (ctx )
7977 if err != nil {
8078 return errors .Wrap (err , "error opening subscription stream" )
8179 }
8280
81+ req := & geyserpb.SubscribeRequest {
82+ Accounts : make (map [string ]* geyserpb.SubscribeRequestFilterAccounts ),
83+ }
84+ req .Accounts ["accounts_subscription" ] = & geyserpb.SubscribeRequestFilterAccounts {
85+ Owner : []string {base58 .Encode (cvm .PROGRAM_ID )},
86+ }
87+ finalizedCommitmentLevel := geyserpb .CommitmentLevel_FINALIZED
88+ req .Commitment = & finalizedCommitmentLevel
89+ err = streamer .Send (req )
90+ if err != nil {
91+ return errors .Wrap (err , "error sending subscription request" )
92+ }
93+
8394 for {
84- update , err := boundedProgramUpdateRecv (ctx , streamer , defaultStreamSubscriptionTimeout )
95+ update , err := boundedRecv (ctx , streamer , defaultStreamSubscriptionTimeout )
8596 if err != nil {
86- return errors .Wrap (err , "error receiving update" )
97+ return errors .Wrap (err , "error recieving update" )
8798 }
8899
89- messageAge := time .Since (update .Ts .AsTime ())
90- if messageAge > defaultStreamSubscriptionTimeout {
91- log .WithField ("message_age" , messageAge ).Warn (ErrSubscriptionFallenBehind .Error ())
92- return ErrSubscriptionFallenBehind
100+ accountUpdate := update .GetAccount ()
101+ if accountUpdate == nil {
102+ continue
93103 }
94104
95105 // Ignore startup updates. We only care about real-time updates due to
96106 // transactions.
97- if update . AccountUpdate .IsStartup {
107+ if accountUpdate .IsStartup {
98108 continue
99109 }
100110
@@ -103,9 +113,42 @@ func (w *Worker) subscribeToProgramUpdatesFromGeyser(ctx context.Context, endpoi
103113 // backing up the Geyser plugin, which kills this subscription and we end up
104114 // missing updates.
105115 select {
106- case w .programUpdatesChan <- update . AccountUpdate :
116+ case w .programUpdatesChan <- accountUpdate :
107117 default :
108118 log .Warn ("dropping update because queue is full" )
109119 }
110120 }
111121}
122+
123+ func newXTokenUnaryClientInterceptor (xToken string ) grpc.UnaryClientInterceptor {
124+ return func (
125+ ctx context.Context ,
126+ method string ,
127+ req , reply interface {},
128+ cc * grpc.ClientConn ,
129+ invoker grpc.UnaryInvoker ,
130+ opts ... grpc.CallOption ,
131+ ) error {
132+ ctx = withXToken (ctx , xToken )
133+ return invoker (ctx , method , req , reply , cc , opts ... )
134+ }
135+ }
136+
137+ func newXTokenStreamClientInterceptor (xToken string ) grpc.StreamClientInterceptor {
138+ return func (
139+ ctx context.Context ,
140+ desc * grpc.StreamDesc ,
141+ cc * grpc.ClientConn ,
142+ method string ,
143+ streamer grpc.Streamer ,
144+ opts ... grpc.CallOption ,
145+ ) (grpc.ClientStream , error ) {
146+ ctx = withXToken (ctx , xToken )
147+ return streamer (ctx , desc , cc , method , opts ... )
148+ }
149+ }
150+
151+ func withXToken (ctx context.Context , xToken string ) context.Context {
152+ md := metadata .Pairs ("x-token" , xToken )
153+ return metadata .NewOutgoingContext (ctx , md )
154+ }
0 commit comments