Fix many instances of contexts and loggers not being propagated (#5340)

This also replaces instances where we manually write "return ctx.Err()"
with "return context.Cause(ctx)" which is functionally identical, but
will also correctly propagate cause errors if present.

pkg/grpc/client.go

@@ -141,7 +141,7 @@ func WaitForReady(ctx context.Context, cc *grpc.ClientConn, timeout time.Duratio
 	for {
 		select {
 		case <-ctx.Done():
-			return ctx.Err()
+			return context.Cause(ctx)
 		case <-ticker.C:
 		}
 

pkg/grpc/databroker/leaser.go

@@ -71,13 +71,13 @@ func (locker *Leaser) Run(ctx context.Context) error {
 		case err == nil:
 			select {
 			case <-ctx.Done():
-				return ctx.Err()
+				return context.Cause(ctx)
 			case <-retryTicker.C:
 			}
 		case errors.Is(err, retryableError{}):
 			select {
 			case <-ctx.Done():
-				return ctx.Err()
+				return context.Cause(ctx)
 			case <-time.After(bo.NextBackOff()):
 			}
 		default:

pkg/grpc/databroker/leaser_test.go

@@ -169,7 +169,7 @@ func TestLeasers(t *testing.T) {
 	fn2 := func(ctx context.Context) error {
 		atomic.AddInt64(&counter, 10)
 		<-ctx.Done()
-		return ctx.Err()
+		return context.Cause(ctx)
 	}
 	leaser := databroker.NewLeasers("TEST", time.Second*30, client, fn1, fn2)
 	err := leaser.Run(context.Background())

pkg/grpc/databroker/reconciler.go

@@ -110,7 +110,7 @@ func (r *Reconciler) reconcileLoop(ctx context.Context) error {
 
 		select {
 		case <-ctx.Done():
-			return ctx.Err()
+			return context.Cause(ctx)
 		case <-r.trigger:
 		}
 	}

pkg/grpc/databroker/sync_test.go

@@ -24,7 +24,7 @@ func Test_SyncLatestRecords(t *testing.T) {
 	defer clearTimeout()
 
 	cc := testutil.NewGRPCServer(t, func(s *grpc.Server) {
-		databrokerpb.RegisterDataBrokerServiceServer(s, databroker.New())
+		databrokerpb.RegisterDataBrokerServiceServer(s, databroker.New(ctx))
 	})
 
 	c := databrokerpb.NewDataBrokerServiceClient(cc)

pkg/grpc/databroker/syncer.go

@@ -72,8 +72,8 @@ type Syncer struct {
 }
 
 // NewSyncer creates a new Syncer.
-func NewSyncer(id string, handler SyncerHandler, options ...SyncerOption) *Syncer {
-	closeCtx, closeCtxCancel := context.WithCancel(context.Background())
+func NewSyncer(ctx context.Context, id string, handler SyncerHandler, options ...SyncerOption) *Syncer {
+	closeCtx, closeCtxCancel := context.WithCancel(context.WithoutCancel(ctx))
 
 	bo := backoff.NewExponentialBackOff()
 	bo.MaxElapsedTime = 0
@@ -120,7 +120,7 @@ func (syncer *Syncer) Run(ctx context.Context) error {
 			log.Ctx(ctx).Error().Err(err).Msg("sync")
 			select {
 			case <-ctx.Done():
-				return ctx.Err()
+				return context.Cause(ctx)
 			case <-time.After(syncer.backoff.NextBackOff()):
 			}
 		}
@@ -133,6 +133,9 @@ func (syncer *Syncer) init(ctx context.Context) error {
 		Type: syncer.cfg.typeURL,
 	})
 	if err != nil {
+		if status.Code(err) == codes.Canceled && ctx.Err() != nil {
+			err = fmt.Errorf("%w: %w", err, context.Cause(ctx))
+		}
 		return fmt.Errorf("error during initial sync: %w", err)
 	}
 	syncer.backoff.Reset()
@@ -167,6 +170,9 @@ func (syncer *Syncer) sync(ctx context.Context) error {
 			syncer.serverVersion = 0
 			return nil
 		} else if err != nil {
+			if status.Code(err) == codes.Canceled && ctx.Err() != nil {
+				err = fmt.Errorf("%w: %w", err, context.Cause(ctx))
+			}
 			return fmt.Errorf("error receiving sync record: %w", err)
 		}
 

pkg/grpc/databroker/syncer_test.go

@@ -157,7 +157,7 @@ func TestSyncer(t *testing.T) {
 
 	clearCh := make(chan struct{})
 	updateCh := make(chan []*Record)
-	syncer := NewSyncer("test", testSyncerHandler{
+	syncer := NewSyncer(ctx, "test", testSyncerHandler{
 		getDataBrokerServiceClient: func() DataBrokerServiceClient {
 			return NewDataBrokerServiceClient(gc)
 		},