use context to close subscription

Author: ysmood

.github/workflows/test.yml

@@ -1,9 +1,13 @@
 name: Test
+
 on:
   push:
 
   pull_request:
 
+env:
+  GODEBUG: tracebackancestors=1000
+
 jobs:
 
   linux:

examples_test.go

@@ -1,17 +1,21 @@
 package goob_test
 
 import (
+	"context"
 	"fmt"
+	"time"
 
 	"github.com/ysmood/goob"
 )
 
 func Example_basic() {
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
 	// create an observable instance
-	ob := goob.New()
-	defer ob.Close()
+	ob := goob.New(ctx)
 
-	events := ob.Subscribe()
+	events := ob.Subscribe(context.TODO())
 
 	// publish events without blocking
 	ob.Publish(1)
@@ -21,10 +25,6 @@ func Example_basic() {
 	// consume events
 	for e := range events {
 		fmt.Print(e)
-
-		if e.(int) == 3 {
-			break
-		}
 	}
 
 	// Output: 123

go.mod

@@ -1,3 +1,5 @@
 module github.com/ysmood/goob
 
 go 1.15
+
+require github.com/ysmood/gotrace v0.6.0

go.sum

@@ -0,0 +1,2 @@
+github.com/ysmood/gotrace v0.6.0 h1:SyI1d4jclswLhg7SWTL6os3L1WOKeNn/ZtzVQF8QmdY=
+github.com/ysmood/gotrace v0.6.0/go.mod h1:TzhIG7nHDry5//eYZDYcTzuJLYQIkykJzCRIo4/dzQM=

goob.go

@@ -1,23 +1,23 @@
 package goob
 
 import (
+	"context"
 	"sync"
 )
 
 // Observable hub
 type Observable struct {
+	ctx         context.Context
 	lock        *sync.Mutex
-	subscribers map[Subscriber]*Pipe
+	subscribers map[Events]func(Event)
 }
 
-// Subscriber type
-type Subscriber <-chan Event
-
 // New observable instance
-func New() *Observable {
+func New(ctx context.Context) *Observable {
 	ob := &Observable{
+		ctx:         ctx,
 		lock:        &sync.Mutex{},
-		subscribers: map[Subscriber]*Pipe{},
+		subscribers: map[Events]func(Event){},
 	}
 	return ob
 }
@@ -27,47 +27,36 @@ func (ob *Observable) Publish(e Event) {
 	ob.lock.Lock()
 	defer ob.lock.Unlock()
 
-	for _, p := range ob.subscribers {
-		p.Write(e)
+	for _, write := range ob.subscribers {
+		write(e)
 	}
 }
 
 // Subscribe message
-func (ob *Observable) Subscribe() Subscriber {
+func (ob *Observable) Subscribe(ctx context.Context) Events {
 	ob.lock.Lock()
 	defer ob.lock.Unlock()
 
-	p := NewPipe()
+	ctx, cancel := context.WithCancel(ctx)
 
-	if ob.subscribers == nil {
-		p.Stop()
-	} else {
-		ob.subscribers[p.Events] = p
-	}
+	write, events := NewPipe(ctx)
 
-	return p.Events
-}
+	ob.subscribers[events] = write
 
-// Unsubscribe from observable
-func (ob *Observable) Unsubscribe(s Subscriber) {
-	ob.lock.Lock()
-	defer ob.lock.Unlock()
-	if p, has := ob.subscribers[s]; has {
-		p.Stop()
-		delete(ob.subscribers, s)
-	}
-}
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-ob.ctx.Done():
+		}
 
-// Close subscribers
-func (ob *Observable) Close() {
-	ob.lock.Lock()
-	defer ob.lock.Unlock()
+		ob.lock.Lock()
+		defer ob.lock.Unlock()
 
-	for _, p := range ob.subscribers {
-		p.Stop()
-	}
+		delete(ob.subscribers, events)
+		cancel()
+	}()
 
-	ob.subscribers = nil
+	return events
 }
 
 // Len of the subscribers

goob_test.go

@@ -1,6 +1,7 @@
 package goob_test
 
 import (
+	"context"
 	"math/rand"
 	"reflect"
 	"runtime"
@@ -10,10 +11,11 @@ import (
 	"time"
 
 	"github.com/ysmood/goob"
+	"github.com/ysmood/gotrace"
 )
 
 func checkLeak(t *testing.T) {
-	// testleak.Check(t, 0)
+	gotrace.CheckLeak(t, 0)
 }
 
 type null struct{}
@@ -27,9 +29,11 @@ func eq(t *testing.T, expected, actual interface{}) {
 func TestNew(t *testing.T) {
 	checkLeak(t)
 
-	ob := goob.New()
-	defer ob.Close()
-	s := ob.Subscribe()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer t.Cleanup(cancel)
+
+	ob := goob.New(ctx)
+	s := ob.Subscribe(ctx)
 	size := 1000
 
 	expected := []int{}
@@ -51,31 +55,31 @@ func TestNew(t *testing.T) {
 	eq(t, expected, result)
 }
 
-func TestUnsubscribe(t *testing.T) {
+func TestCancel(t *testing.T) {
 	checkLeak(t)
 
-	ob := goob.New()
-
-	s := ob.Subscribe()
-	ob.Unsubscribe(s)
+	ob := goob.New(context.Background())
 
+	ctx, cancel := context.WithCancel(context.Background())
+	ob.Subscribe(ctx)
+	cancel()
 	time.Sleep(10 * time.Millisecond)
-
 	eq(t, ob.Len(), 0)
 }
 
 func TestClosed(t *testing.T) {
 	checkLeak(t)
 
-	ob := goob.New()
-	ob.Subscribe()
-	ob.Close()
+	ctx, cancel := context.WithCancel(context.Background())
 
-	s := ob.Subscribe()
+	ob := goob.New(ctx)
+	ob.Subscribe(ctx)
+	cancel()
+
+	s := ob.Subscribe(context.Background())
 	_, ok := <-s
 
 	ob.Publish(1)
-	ob.Unsubscribe(s)
 
 	eq(t, ok, false)
 	eq(t, ob.Len(), 0)
@@ -84,11 +88,13 @@ func TestClosed(t *testing.T) {
 func TestMultipleConsumers(t *testing.T) {
 	checkLeak(t)
 
-	ob := goob.New()
-	defer ob.Close()
-	s1 := ob.Subscribe()
-	s2 := ob.Subscribe()
-	s3 := ob.Subscribe()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer t.Cleanup(cancel)
+
+	ob := goob.New(ctx)
+	s1 := ob.Subscribe(ctx)
+	s2 := ob.Subscribe(ctx)
+	s3 := ob.Subscribe(ctx)
 	size := 1000
 
 	expected := []int{}
@@ -136,9 +142,11 @@ func TestMultipleConsumers(t *testing.T) {
 func TestSlowConsumer(t *testing.T) {
 	checkLeak(t)
 
-	ob := goob.New()
-	defer ob.Close()
-	s := ob.Subscribe()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer t.Cleanup(cancel)
+
+	ob := goob.New(ctx)
+	s := ob.Subscribe(ctx)
 
 	ob.Publish(1)
 
@@ -160,9 +168,11 @@ func TestMonkey(t *testing.T) {
 	run := func() {
 		defer wg.Done()
 
-		ob := goob.New()
-		defer ob.Close()
-		s := ob.Subscribe()
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		ob := goob.New(ctx)
+		s := ob.Subscribe(ctx)
 
 		go func() {
 			for range make([]null, size) {
@@ -198,9 +208,11 @@ func TestMonkey(t *testing.T) {
 }
 
 func BenchmarkPublish(b *testing.B) {
-	ob := goob.New()
-	defer ob.Close()
-	s := ob.Subscribe()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer b.Cleanup(cancel)
+
+	ob := goob.New(ctx)
+	s := ob.Subscribe(ctx)
 
 	for i := 0; i < runtime.NumCPU(); i++ {
 		go func() {
@@ -219,9 +231,11 @@ func BenchmarkPublish(b *testing.B) {
 }
 
 func BenchmarkConsume(b *testing.B) {
-	ob := goob.New()
-	defer ob.Close()
-	s := ob.Subscribe()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer b.Cleanup(cancel)
+
+	ob := goob.New(ctx)
+	s := ob.Subscribe(ctx)
 
 	for i := 0; i < b.N; i++ {
 		ob.Publish(nil)

pipe.go

@@ -1,27 +1,23 @@
 package goob
 
 import (
+	"context"
 	"sync"
 )
 
 // Event interface
 type Event interface{}
 
-// Pipe the Event via Write to Events. Events uses an internal buffer so it won't block Write.
-// Call Stop to abort.
-type Pipe struct {
-	Write  func(Event)
-	Events <-chan Event
-	Stop   func()
-}
+// Events channel
+type Events <-chan Event
 
-// NewPipe instance
-func NewPipe() *Pipe {
+// NewPipe instance.
+// Pipe the Event via Write to Events. Events uses an internal buffer so it won't block Write.
+func NewPipe(ctx context.Context) (Write func(Event), Events <-chan Event) {
 	events := make(chan Event)
 	lock := sync.Mutex{}
 	buf := []Event{} // using slice is faster than linked-list in general cases
 	wait := make(chan struct{}, 1)
-	stop := make(chan struct{})
 
 	write := func(e Event) {
 		lock.Lock()
@@ -31,7 +27,7 @@ func NewPipe() *Pipe {
 
 		if len(wait) == 0 {
 			select {
-			case <-stop:
+			case <-ctx.Done():
 				return
 			case wait <- struct{}{}:
 			}
@@ -49,19 +45,19 @@ func NewPipe() *Pipe {
 
 			for _, e := range section {
 				select {
-				case <-stop:
+				case <-ctx.Done():
 					return
 				case events <- e:
 				}
 			}
 
 			select {
-			case <-stop:
+			case <-ctx.Done():
 				return
 			case <-wait:
 			}
 		}
 	}()
 
-	return &Pipe{write, events, func() { close(stop) }}
+	return write, events
 }