Use struct for cancellation callbacks

Author: cschleiden

internal/sync/channel.go

@@ -12,14 +12,20 @@ type Channel[T any] interface {
 	Close()
 }
 
+type Receiver[T any] struct {
+	Receive func(v T, ok bool)
+}
+
 type ChannelInternal[T any] interface {
 	Closed() bool
 
 	ReceiveNonBlocking() (v T, ok bool)
 
 	// AddReceiveCallback adds a callback that is called once when a value is sent to the channel. This is similar
 	// to the blocking `Receive` method, but is not blocking a coroutine.
-	AddReceiveCallback(cb func(v T, ok bool))
+	AddReceiveCallback(rcb *Receiver[T])
+
+	RemoveReceiveCallback(rcb *Receiver[T])
 }
 
 // Ensure channel implementation support internal interface
@@ -40,7 +46,7 @@ func NewBufferedChannel[T any](size int) Channel[T] {
 
 type channel[T any] struct {
 	c         []T
-	receivers []func(value T, ok bool)
+	receivers []*Receiver[T]
 	senders   []func() T
 	closed    bool
 	size      int
@@ -62,7 +68,7 @@ func (c *channel[T]) Close() {
 
 		// Send zero value to pending receiver
 		var v T
-		r(v, false)
+		r.Receive(v, false)
 	}
 }
 
@@ -119,10 +125,12 @@ func (c *channel[T]) Receive(ctx Context) (v T, ok bool) {
 
 		// Register handler to receive value once
 		if !addedListener {
-			cb := func(rv T, rok bool) {
-				receivedValue = true
-				v = rv
-				ok = rok
+			cb := &Receiver[T]{
+				Receive: func(rv T, rok bool) {
+					receivedValue = true
+					v = rv
+					ok = rok
+				},
 			}
 
 			c.receivers = append(c.receivers, cb)
@@ -176,7 +184,7 @@ func (c *channel[T]) trySend(v T) bool {
 		c.receivers[0] = nil
 		c.receivers = c.receivers[1:]
 
-		r(v, true)
+		r.Receive(v, true)
 
 		return true
 	}
@@ -223,10 +231,20 @@ func (c *channel[T]) hasCapacity() bool {
 	return len(c.c) < c.size
 }
 
-func (c *channel[T]) AddReceiveCallback(cb func(v T, ok bool)) {
+func (c *channel[T]) AddReceiveCallback(cb *Receiver[T]) {
 	c.receivers = append(c.receivers, cb)
 }
 
+func (c *channel[T]) RemoveReceiveCallback(cb *Receiver[T]) {
+	for i, r := range c.receivers {
+		if r == cb {
+			c.receivers[i] = nil
+			c.receivers = append(c.receivers[:i], c.receivers[i+1:]...)
+			return
+		}
+	}
+}
+
 func (c *channel[T]) Closed() bool {
 	return c.closed
 }

internal/sync/channel_test.go

@@ -413,3 +413,60 @@ func Test_Channel_Buffered(t *testing.T) {
 		})
 	}
 }
+
+func Test_CancellationHandler_Add(t *testing.T) {
+	ctx, cancel := WithCancel(Background())
+
+	f := 1
+
+	c := NewChannel[int]()
+	ic := c.(*channel[int])
+
+	cr := NewCoroutine(ctx, func(ctx Context) error {
+		r := &Receiver[int]{
+			Receive: func(_ int, ok bool) {
+				f++
+			},
+		}
+
+		ic.AddReceiveCallback(r)
+
+		c.Send(ctx, 42)
+		return nil
+	})
+
+	cr.Execute()
+
+	cancel()
+
+	require.Equal(t, 2, f)
+}
+
+func Test_CancellationHandler_Remove(t *testing.T) {
+	ctx, cancel := WithCancel(Background())
+
+	f := 1
+
+	c := NewChannel[int]()
+	ic := c.(*channel[int])
+
+	cr := NewCoroutine(ctx, func(ctx Context) error {
+		r := &Receiver[int]{
+			Receive: func(_ int, ok bool) {
+				f++
+			},
+		}
+
+		ic.AddReceiveCallback(r)
+		ic.RemoveReceiveCallback(r)
+
+		c.Send(ctx, 42)
+		return nil
+	})
+
+	cr.Execute()
+
+	cancel()
+
+	require.Equal(t, 1, f)
+}