@@ -124,34 +124,28 @@ func (bkr *GRPCBroker) Publish(ctx context.Context, pubReq *pbv1.PublishRequest)
124124 return & emptypb.Empty {}, nil
125125}
126126
127- // register registers a subscriber and return client id and error channel.
128- func (bkr * GRPCBroker ) register (ctx context.Context ,
127+ // registerSubscriber registers a subscriber with a pre-generated ID.
128+ // The subscription header must already be sent before calling this function.
129+ func (bkr * GRPCBroker ) registerSubscriber (ctx context.Context ,
130+ id string ,
129131 dataType types.CloudEventsDataType ,
130132 subReq * pbv1.SubscriptionRequest ,
131- subServer pbv1.CloudEventService_SubscribeServer ,
132- handler resourceHandler ) (string , error ) {
133+ handler resourceHandler ) error {
133134 logger := klog .FromContext (ctx )
134135
135136 bkr .mu .Lock ()
136137 defer bkr .mu .Unlock ()
137138
138- id := uuid .NewString ()
139+ logger .Info ("registering subscriber" , "id" , id , "clusterName" , subReq .ClusterName , "dataType" , dataType )
140+
139141 bkr .subscribers [id ] = & subscriber {
140142 clusterName : subReq .ClusterName ,
141143 dataType : dataType ,
142144 handler : handler ,
143145 }
144146
145- // Signal subscriber is registered
146- if err := subServer .SendHeader (metadata .Pairs (constants .GRPCSubscriptionIDKey , id )); err != nil {
147- logger .Error (err , "failed to send subscription header, unregister subscriber" , "subID" , id )
148- delete (bkr .subscribers , id )
149- return "" , err
150- }
151- logger .V (4 ).Info ("register a subscriber" , "id" , id , "clusterName" , subReq .ClusterName , "dataType" , dataType )
152147 metrics .IncGRPCCESubscribersMetric (subReq .ClusterName , dataType .String ())
153-
154- return id , nil
148+ return nil
155149}
156150
157151// unregister a subscriber by id
@@ -181,10 +175,17 @@ func (bkr *GRPCBroker) Subscribe(subReq *pbv1.SubscriptionRequest, subServer pbv
181175 return fmt .Errorf ("invalid subscription request: invalid data type %v" , err )
182176 }
183177
178+ // Generate subscription ID and send header IMMEDIATELY, before any other operations
179+ // This ensures the client receives the header as soon as possible after the stream is established
180+ subID := uuid .NewString ()
181+ if err := subServer .SendHeader (metadata .Pairs (constants .GRPCSubscriptionIDKey , subID )); err != nil {
182+ return fmt .Errorf ("failed to send subscription header for subID %s: %w" , subID , err )
183+ }
184+
184185 subCtx , cancel := context .WithCancel (subServer .Context ())
185186 defer cancel ()
186187
187- logger := klog .FromContext (subCtx ).WithValues ("clusterName" , subReq .ClusterName )
188+ logger := klog .FromContext (subCtx ).WithValues ("clusterName" , subReq .ClusterName , "subID" , subID )
188189
189190 // TODO make the channel size configurable
190191 eventCh := make (chan * pbv1.CloudEvent , 100 )
@@ -195,6 +196,35 @@ func (bkr *GRPCBroker) Subscribe(subReq *pbv1.SubscriptionRequest, subServer pbv
195196 }
196197 sendErrCh := make (chan error , 1 )
197198
199+ // Register the subscriber with the ID we already created and sent in the header
200+ err = bkr .registerSubscriber (klog .NewContext (subCtx , logger ), subID , * dataType , subReq , func (handlerCtx context.Context , subID string , evt * cloudevents.Event ) error {
201+ // convert the cloudevents.Event to pbv1.CloudEvent
202+ // WARNING: don't use "pbEvt, err := pb.ToProto(evt)" to convert cloudevent to protobuf
203+ pbEvt := & pbv1.CloudEvent {}
204+ if err := grpcprotocol .WritePBMessage (handlerCtx , binding .ToMessage (evt ), pbEvt ); err != nil {
205+ return fmt .Errorf ("failed to convert cloudevent to protobuf for resource(%s): %v" , evt .ID (), err )
206+ }
207+
208+ // send the cloudevent to the subscriber
209+ klog .FromContext (handlerCtx ).V (4 ).Info ("sending the event to spec subscribers" ,
210+ "subID" , subID , "eventType" , evt .Type (), "extensions" , evt .Extensions ())
211+ select {
212+ case eventCh <- pbEvt :
213+ case <- subCtx .Done ():
214+ // The context of the stream has been canceled or completed.
215+ // This could happen if:
216+ // - The client closed the connection or canceled the stream.
217+ // - The server closed the stream, potentially due to a shutdown.
218+ // No error is returned here because the stream closure is expected.
219+ return nil
220+ }
221+
222+ return nil
223+ })
224+ if err != nil {
225+ return err
226+ }
227+
198228 // send events
199229 // The grpc send is not concurrency safe and non-blocking, see: https://github.com/grpc/grpc-go/blob/v1.75.1/stream.go#L1571
200230 // Return the error without wrapping, as it includes the gRPC error code and message for further handling.
@@ -238,34 +268,6 @@ func (bkr *GRPCBroker) Subscribe(subReq *pbv1.SubscriptionRequest, subServer pbv
238268 }
239269 }()
240270
241- subID , err := bkr .register (klog .NewContext (subCtx , logger ), * dataType , subReq , subServer , func (handlerCtx context.Context , subID string , evt * cloudevents.Event ) error {
242- // convert the cloudevents.Event to pbv1.CloudEvent
243- // WARNING: don't use "pbEvt, err := pb.ToProto(evt)" to convert cloudevent to protobuf
244- pbEvt := & pbv1.CloudEvent {}
245- if err := grpcprotocol .WritePBMessage (handlerCtx , binding .ToMessage (evt ), pbEvt ); err != nil {
246- return fmt .Errorf ("failed to convert cloudevent to protobuf for resource(%s): %v" , evt .ID (), err )
247- }
248-
249- // send the cloudevent to the subscriber
250- logger .V (4 ).Info ("sending the event to spec subscribers" ,
251- "subID" , subID , "eventType" , evt .Type (), "extensions" , evt .Extensions ())
252- select {
253- case eventCh <- pbEvt :
254- case <- subCtx .Done ():
255- // The context of the stream has been canceled or completed.
256- // This could happen if:
257- // - The client closed the connection or canceled the stream.
258- // - The server closed the stream, potentially due to a shutdown.
259- // No error is returned here because the stream closure is expected.
260- return nil
261- }
262-
263- return nil
264- })
265- if err != nil {
266- return err
267- }
268-
269271 if heartbeater != nil {
270272 go heartbeater .Start (subCtx )
271273 }
0 commit comments