feat: add context management to non buffered channels

Author: blockchainluffy

rolling-shutter/keyperimpl/gnosis/keyper.go

@@ -295,17 +295,29 @@ func (kpr *Keyper) processInputs(ctx context.Context) error {
 	}
 }
 
-func (kpr *Keyper) channelNewBlock(_ context.Context, ev *syncevent.LatestBlock) error {
-	kpr.newBlocks <- ev
-	return nil
+func (kpr *Keyper) channelNewBlock(ctx context.Context, ev *syncevent.LatestBlock) error {
+	select {
+	case kpr.newBlocks <- ev:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
-func (kpr *Keyper) channelNewKeyperSet(_ context.Context, ev *syncevent.KeyperSet) error {
-	kpr.newKeyperSets <- ev
-	return nil
+func (kpr *Keyper) channelNewKeyperSet(ctx context.Context, ev *syncevent.KeyperSet) error {
+	select {
+	case kpr.newKeyperSets <- ev:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
-func (kpr *Keyper) channelNewEonPublicKey(_ context.Context, key keyper.EonPublicKey) error {
-	kpr.newEonPublicKeys <- key
-	return nil
+func (kpr *Keyper) channelNewEonPublicKey(ctx context.Context, key keyper.EonPublicKey) error {
+	select {
+	case kpr.newEonPublicKeys <- key:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }

rolling-shutter/keyperimpl/gnosis/newslot.go

@@ -266,7 +266,12 @@ func (kpr *Keyper) triggerDecryption(
 		Int("num-identities", len(trigger.IdentityPreimages)).
 		Int64("tx-pointer", txPointer).
 		Msg("sending decryption trigger")
-	kpr.decryptionTriggerChannel <- event
+
+	select {
+	case kpr.decryptionTriggerChannel <- event:
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 
 	return nil
 }

rolling-shutter/keyperimpl/shutterservice/keyper.go

@@ -274,17 +274,29 @@ func (kpr *Keyper) processInputs(ctx context.Context) error {
 	}
 }
 
-func (kpr *Keyper) channelNewEonPublicKey(_ context.Context, key keyper.EonPublicKey) error {
-	kpr.newEonPublicKeys <- key
-	return nil
+func (kpr *Keyper) channelNewEonPublicKey(ctx context.Context, key keyper.EonPublicKey) error {
+	select {
+	case kpr.newEonPublicKeys <- key:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
-func (kpr *Keyper) channelNewBlock(_ context.Context, ev *syncevent.LatestBlock) error {
-	kpr.newBlocks <- ev
-	return nil
+func (kpr *Keyper) channelNewBlock(ctx context.Context, ev *syncevent.LatestBlock) error {
+	select {
+	case kpr.newBlocks <- ev:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }
 
-func (kpr *Keyper) channelNewKeyperSet(_ context.Context, ev *syncevent.KeyperSet) error {
-	kpr.newKeyperSets <- ev
-	return nil
+func (kpr *Keyper) channelNewKeyperSet(ctx context.Context, ev *syncevent.KeyperSet) error {
+	select {
+	case kpr.newKeyperSets <- ev:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }

rolling-shutter/medley/channel/channel.go

@@ -12,7 +12,11 @@ func Forward[T any](ctx context.Context,
 			if !ok {
 				return nil
 			}
-			send <- val
+			select {
+			case send <- val:
+			case <-ctx.Done():
+				return ctx.Err()
+			}
 		case <-ctx.Done():
 			return ctx.Err()
 		}

rolling-shutter/medley/slotticker/slotticker.go

@@ -92,7 +92,11 @@ func (t *SlotTicker) run(ctx context.Context) error {
 
 		timeToNextSlot := nextTickTime.Sub(now)
 		timer.Reset(timeToNextSlot)
-		<-timer.C
+		select {
+		case <-timer.C:
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 
 		if err := t.tick(ctx, nextSlotNumber); err != nil {
 			return err