@@ -266,7 +266,6 @@ func (c *CfClient) streamConnect(ctx context.Context) {
266266 // while this is happening we set streamConnectedBool to true - if any errors happen
267267 // in this process streamConnectedBool will be set back to false by the streamDisconnected function
268268 conn .Connect (ctx , c .environmentID , c .sdkKey )
269- c .streamConnectedBool = true
270269}
271270
272271func (c * CfClient ) initAuthentication (ctx context.Context ) error {
@@ -390,6 +389,31 @@ func (c *CfClient) authenticate(ctx context.Context) error {
390389 if bearerTokenProviderErr != nil {
391390 return bearerTokenProviderErr
392391 }
392+
393+ // Use a custom transport which adds headers for tracking usage
394+ // The `WithRequestEditorFn` cannot be used for SSE requests, so we need to provide a custom transport to the
395+ // http client so that these headers can be added to all requests.
396+ getHeadersFn := func () (map [string ]string , error ) {
397+ return map [string ]string {
398+ "User-Agent" : "GoSDK/" + analyticsservice .SdkVersion ,
399+ "Harness-SDK-Info" : fmt .Sprintf ("Go %s Server" , analyticsservice .SdkVersion ),
400+ "Harness-EnvironmentID" : c .environmentID ,
401+ }, nil
402+ }
403+
404+ // Wrap the httpClient's transport with our own custom transport, which currently just adds extra headers
405+ // for analytics purposes.
406+ // If the httpClient doesn't have a Transport we can honour, then just use a default transport.
407+ var baseTransport http.RoundTripper
408+ if c .config .httpClient .Transport != nil {
409+ baseTransport = c .config .httpClient .Transport
410+ } else {
411+ baseTransport = http .DefaultTransport
412+ }
413+ customTrans := NewCustomTransport (baseTransport , getHeadersFn )
414+
415+ c .config .httpClient .Transport = customTrans
416+
393417 restClient , err := rest .NewClientWithResponses (c .config .url ,
394418 rest .WithRequestEditorFn (bearerTokenProvider .Intercept ),
395419 rest .WithRequestEditorFn (c .InterceptAddCluster ),
@@ -425,70 +449,68 @@ func (c *CfClient) stream(ctx context.Context) {
425449 c .streamConnect (ctx )
426450
427451 streamingRetryStrategy := c .config .streamingRetryStrategy
428- // Create an initial ticker to handle the case where we don't open a succesful connection on our first attempt
429- ticker := backoff .NewTicker (streamingRetryStrategy )
430452
431- defer ticker .Stop ()
432453 reconnectionAttempt := 1
454+
433455 for {
434456 select {
435457 case <- ctx .Done ():
436458 c .config .Logger .Infof ("%s Stream stopped" , sdk_codes .StreamStop )
437- if ticker != nil {
438- ticker .Stop ()
439- }
440459 return
441460
442461 case <- c .streamConnectedChan :
443462 c .config .Logger .Infof ("%s Stream successfully connected" , sdk_codes .StreamStarted )
444463 c .config .Logger .Infof ("%s Polling Stopped" , sdk_codes .PollStop )
445464
446- // Reset the reconnection attempt and ticker
447- reconnectionAttempt = 1
465+ // Ensure reconnection strategy is reset
448466 streamingRetryStrategy .Reset ()
449- ticker .Stop ()
450- ticker = nil
467+ reconnectionAttempt = 1
451468
452- case err := <- c .streamDisconnectedChan :
453- c .config .Logger .Warnf ("%s Stream disconnected: %s" , sdk_codes .StreamDisconnected , err )
454- c .config .Logger .Infof ("%s Polling started, interval: %v seconds" , sdk_codes .PollStart , c .config .pullInterval )
455469 c .mux .RLock ()
456- c .streamConnectedBool = false
470+ c .streamConnectedBool = true
457471 c .mux .RUnlock ()
458472
459- // If an eventStreamListener has been passed to the Proxy lets notify it of the disconnected
460- // to let it know something is up with the stream it has been listening to
461- if c .config .eventStreamListener != nil {
462- c .config .eventStreamListener .Pub (context .Background (), stream.Event {
463- APIKey : c .sdkKey ,
464- Environment : c .environmentID ,
465- Err : stream .ErrStreamDisconnect ,
466- })
467- }
468-
469- if ticker == nil {
470- ticker = backoff .NewTicker (streamingRetryStrategy )
471- }
472-
473- // Note, the retry interval logged here is just an approximate value.
474- // NextBackOff() will not be exactly the same value when used by the underlying ticker.
475- // There is currently no way to get the interval that the ticker has used.
476- c .config .Logger .Infof ("%s Retrying stream connection in %fs (attempt %d)" , sdk_codes .StreamRetry , streamingRetryStrategy .NextBackOff ().Seconds (), reconnectionAttempt )
473+ case err := <- c .streamDisconnectedChan :
474+ c .notifyStreamDisconnect (err )
475+
476+ nextBackOff := streamingRetryStrategy .NextBackOff ()
477+ c .config .Logger .Infof ("%s Retrying stream connection in %fs (attempt %d)" , sdk_codes .StreamRetry , nextBackOff .Seconds (), reconnectionAttempt )
478+ c .handleStreamDisconnect (ctx , nextBackOff )
479+
477480 reconnectionAttempt += 1
478481
479- // Backoff before retrying
480- select {
481- case <- ticker .C :
482- case <- ctx .Done ():
483- c .config .Logger .Infof ("%s Stream stopped during reconnection" , sdk_codes .StreamStop )
484- ticker .Stop ()
485- return
486- }
487- c .streamConnect (ctx )
488482 }
489483 }
490484}
491485
486+ func (c * CfClient ) handleStreamDisconnect (ctx context.Context , nextBackOff time.Duration ) {
487+ select {
488+ case <- time .After (nextBackOff ):
489+ c .streamConnect (ctx )
490+ case <- ctx .Done ():
491+ // Context was cancelled, stop trying to reconnect
492+ c .config .Logger .Infof ("%s Stream stopped during reconnection" , sdk_codes .StreamStop )
493+ return
494+ }
495+ }
496+
497+ func (c * CfClient ) notifyStreamDisconnect (err error ) {
498+ c .mux .RLock ()
499+ c .streamConnectedBool = false
500+ c .mux .RUnlock ()
501+ // If an eventStreamListener has been passed to the Proxy lets notify it of the disconnected
502+ // to let it know something is up with the stream it has been listening to
503+ if c .config .eventStreamListener != nil {
504+ c .config .eventStreamListener .Pub (context .Background (), stream.Event {
505+ APIKey : c .sdkKey ,
506+ Environment : c .environmentID ,
507+ Err : stream .ErrStreamDisconnect ,
508+ })
509+ }
510+ c .config .Logger .Warnf ("%s Stream disconnected: %s" , sdk_codes .StreamDisconnected , err )
511+ c .config .Logger .Infof ("%s Polling started, interval: %v seconds" , sdk_codes .PollStart , c .config .pullInterval )
512+ }
513+
492514func (c * CfClient ) pullCronJob (ctx context.Context ) {
493515 poll := func () {
494516 c .mux .RLock ()
0 commit comments