Repository: looplab/eventhorizon
Pull request #342 (2021-10-04 14:02:13)

Fix / Add closers for graceful shutdown
### Description

- Add Close() to EventBus, remove Wait(). Previous behaviour of cancelling the handling context no longer works.
- Add Close() to EventStore.
- Add Close() to ReadRepo.

### Affected Components

- Event Bus
- Event Store
- Read Repo

### Related Issues

#123

### Solution and Design

Short description of the solution.

### Steps to test and verify

1. ...



Patch:
diff --git a/eventbus.go b/eventbus.go
index 3fca2dd..57b289c 100644
--- a/eventbus.go
+++ b/eventbus.go
@@ -34,8 +34,8 @@ type EventBus interface {
 	// Errors returns an error channel where async handling errors are sent.
 	Errors() <-chan EventBusError
 
-	// Wait wait for all handlers to be cancelled by their context.
-	Wait()
+	// Close closes the EventBus and waits for all handlers to finish.
+	Close() error
 }
 
 // EventBusError is an async error containing the error returned from a handler
diff --git a/eventbus/gcp/eventbus.go b/eventbus/gcp/eventbus.go
index 916bb1b..1c1949b 100644
--- a/eventbus/gcp/eventbus.go
+++ b/eventbus/gcp/eventbus.go
@@ -39,16 +39,22 @@ type EventBus struct {
 	registered   map[eh.EventHandlerType]struct{}
 	registeredMu sync.RWMutex
 	errCh        chan eh.EventBusError
+	cctx         context.Context
+	cancel       context.CancelFunc
 	wg           sync.WaitGroup
 	codec        eh.EventCodec
 }
 
 // NewEventBus creates an EventBus, with optional GCP connection settings.
 func NewEventBus(projectID, appID string, options ...Option) (*EventBus, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+
 	b := &EventBus{
 		appID:      appID,
 		registered: map[eh.EventHandlerType]struct{}{},
 		errCh:      make(chan eh.EventBusError, 100),
+		cctx:       ctx,
+		cancel:     cancel,
 		codec:      &json.EventCodec{},
 	}
 
@@ -63,9 +69,8 @@ func NewEventBus(projectID, appID string, options ...Option) (*EventBus, error)
 	}
 
 	// Create the GCP pubsub client.
-	ctx := context.Background()
 	var err error
-	b.client, err = pubsub.NewClient(ctx, projectID, b.clientOpts...)
+	b.client, err = pubsub.NewClient(b.cctx, projectID, b.clientOpts...)
 	if err != nil {
 		return nil, err
 	}
@@ -73,10 +78,10 @@ func NewEventBus(projectID, appID string, options ...Option) (*EventBus, error)
 	// Get or create the topic.
 	name := appID + "_events"
 	b.topic = b.client.Topic(name)
-	if ok, err := b.topic.Exists(ctx); err != nil {
+	if ok, err := b.topic.Exists(b.cctx); err != nil {
 		return nil, err
 	} else if !ok {
-		if b.topic, err = b.client.CreateTopic(ctx, name); err != nil {
+		if b.topic, err = b.client.CreateTopic(b.cctx, name); err != nil {
 			return nil, err
 		}
 	}
@@ -195,8 +200,7 @@ func (b *EventBus) AddHandler(ctx context.Context, m eh.EventMatcher, h eh.Event
 	b.registered[h.HandlerType()] = struct{}{}
 
 	// Handle until context is cancelled.
-	b.wg.Add(1)
-	go b.handle(ctx, m, h, sub)
+	go b.handle(m, h, sub)
 
 	return nil
 }
@@ -206,20 +210,28 @@ func (b *EventBus) Errors() <-chan eh.EventBusError {
 	return b.errCh
 }
 
-// Wait for all channels to close in the event bus group
-func (b *EventBus) Wait() {
+// Close implements the Close method of the eventhorizon.EventBus interface.
+func (b *EventBus) Close() error {
+	// Stop publishing.
+	b.topic.Stop()
+
+	// Stop handling.
+	b.cancel()
 	b.wg.Wait()
+
+	return b.client.Close()
 }
 
 // Handles all events coming in on the channel.
-func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHandler, sub *pubsub.Subscription) {
+func (b *EventBus) handle(m eh.EventMatcher, h eh.EventHandler, sub *pubsub.Subscription) {
+	b.wg.Add(1)
 	defer b.wg.Done()
 
 	for {
-		if err := sub.Receive(ctx, b.handler(m, h)); err != nil {
+		if err := sub.Receive(b.cctx, b.handler(m, h)); err != nil {
 			err = fmt.Errorf("could not receive: %w", err)
 			select {
-			case b.errCh <- eh.EventBusError{Err: err, Ctx: ctx}:
+			case b.errCh <- eh.EventBusError{Err: err}:
 			default:
 				log.Printf("eventhorizon: missed error in GCP event bus: %s", err)
 			}
@@ -227,6 +239,7 @@ func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHand
 			time.Sleep(time.Second)
 			continue
 		}
+
 		return
 	}
 }
diff --git a/eventbus/kafka/eventbus.go b/eventbus/kafka/eventbus.go
index 25108b6..907c41e 100644
--- a/eventbus/kafka/eventbus.go
+++ b/eventbus/kafka/eventbus.go
@@ -40,19 +40,24 @@ type EventBus struct {
 	registered   map[eh.EventHandlerType]struct{}
 	registeredMu sync.RWMutex
 	errCh        chan eh.EventBusError
+	cctx         context.Context
+	cancel       context.CancelFunc
 	wg           sync.WaitGroup
 	codec        eh.EventCodec
 }
 
 // NewEventBus creates an EventBus, with optional GCP connection settings.
 func NewEventBus(addr, appID string, options ...Option) (*EventBus, error) {
-	topic := appID + "_events"
+	ctx, cancel := context.WithCancel(context.Background())
+
 	b := &EventBus{
 		addr:       addr,
 		appID:      appID,
-		topic:      topic,
+		topic:      appID + "_events",
 		registered: map[eh.EventHandlerType]struct{}{},
 		errCh:      make(chan eh.EventBusError, 100),
+		cctx:       ctx,
+		cancel:     cancel,
 		codec:      &json.EventCodec{},
 	}
 
@@ -75,7 +80,7 @@ func NewEventBus(addr, appID string, options ...Option) (*EventBus, error) {
 	for i := 0; i < 10; i++ {
 		resp, err = client.CreateTopics(context.Background(), &kafka.CreateTopicsRequest{
 			Topics: []kafka.TopicConfig{{
-				Topic:             topic,
+				Topic:             b.topic,
 				NumPartitions:     5,
 				ReplicationFactor: 1,
 			}},
@@ -92,7 +97,7 @@ func NewEventBus(addr, appID string, options ...Option) (*EventBus, error) {
 	if resp == nil {
 		return nil, fmt.Errorf("could not get/create Kafka topic in time: %w", err)
 	}
-	if topicErr, ok := resp.Errors[topic]; ok && topicErr != nil {
+	if topicErr, ok := resp.Errors[b.topic]; ok && topicErr != nil {
 		if !errors.Is(topicErr, kafka.TopicAlreadyExists) {
 			return nil, fmt.Errorf("invalid Kafka topic: %w", topicErr)
 		}
@@ -100,7 +105,7 @@ func NewEventBus(addr, appID string, options ...Option) (*EventBus, error) {
 
 	b.writer = &kafka.Writer{
 		Addr:         kafka.TCP(addr),
-		Topic:        topic,
+		Topic:        b.topic,
 		BatchSize:    1,                // Write every event to the bus without delay.
 		RequiredAcks: kafka.RequireOne, // Stronger consistency.
 		Balancer:     &kafka.Hash{},    // Hash by aggregate ID.
@@ -206,8 +211,7 @@ func (b *EventBus) AddHandler(ctx context.Context, m eh.EventMatcher, h eh.Event
 	b.registered[h.HandlerType()] = struct{}{}
 
 	// Handle until context is cancelled.
-	b.wg.Add(1)
-	go b.handle(ctx, m, h, r)
+	go b.handle(m, h, r)
 
 	return nil
 }
@@ -217,28 +221,29 @@ func (b *EventBus) Errors() <-chan eh.EventBusError {
 	return b.errCh
 }
 
-// Wait for all channels to close in the event bus group
-func (b *EventBus) Wait() {
+// Close implements the Close method of the eventhorizon.EventBus interface.
+func (b *EventBus) Close() error {
+	b.cancel()
 	b.wg.Wait()
-	if err := b.writer.Close(); err != nil {
-		log.Printf("eventhorizon: failed to close Kafka writer: %s", err)
-	}
+
+	return b.writer.Close()
 }
 
 // Handles all events coming in on the channel.
-func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHandler, r *kafka.Reader) {
+func (b *EventBus) handle(m eh.EventMatcher, h eh.EventHandler, r *kafka.Reader) {
+	b.wg.Add(1)
 	defer b.wg.Done()
+
 	handler := b.handler(m, h, r)
 
 	for {
-		msg, err := r.FetchMessage(ctx)
+		msg, err := r.FetchMessage(b.cctx)
 		if errors.Is(err, context.Canceled) {
 			break
-		}
-		if err != nil {
+		} else if err != nil {
 			err = fmt.Errorf("could not fetch message: %w", err)
 			select {
-			case b.errCh <- eh.EventBusError{Err: err, Ctx: ctx}:
+			case b.errCh <- eh.EventBusError{Err: err}:
 			default:
 				log.Printf("eventhorizon: missed error in Kafka event bus: %s", err)
 			}
@@ -248,17 +253,17 @@ func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHand
 		}
 
 		var noBusError eh.EventBusError
-		if err := handler(ctx, msg); err != noBusError {
+		if err := handler(b.cctx, msg); err != noBusError {
 			select {
 			case b.errCh <- err:
 			default:
 				log.Printf("eventhorizon: missed error in Kafka event bus: %s", err)
 			}
 		} else {
-			if err := r.CommitMessages(ctx, msg); err != nil {
+			if err := r.CommitMessages(b.cctx, msg); err != nil {
 				err = fmt.Errorf("could not commit message: %w", err)
 				select {
-				case b.errCh <- eh.EventBusError{Err: err, Ctx: ctx}:
+				case b.errCh <- eh.EventBusError{Err: err}:
 				default:
 					log.Printf("eventhorizon: missed error in Kafka event bus: %s", err)
 				}
diff --git a/eventbus/local/eventbus.go b/eventbus/local/eventbus.go
index 1b21cc6..3f3c2c4 100644
--- a/eventbus/local/eventbus.go
+++ b/eventbus/local/eventbus.go
@@ -35,16 +35,22 @@ type EventBus struct {
 	registered   map[eh.EventHandlerType]struct{}
 	registeredMu sync.RWMutex
 	errCh        chan eh.EventBusError
+	cctx         context.Context
+	cancel       context.CancelFunc
 	wg           sync.WaitGroup
 	codec        eh.EventCodec
 }
 
 // NewEventBus creates a EventBus.
 func NewEventBus(options ...Option) *EventBus {
+	ctx, cancel := context.WithCancel(context.Background())
+
 	b := &EventBus{
 		group:      NewGroup(),
 		registered: map[eh.EventHandlerType]struct{}{},
 		errCh:      make(chan eh.EventBusError, 100),
+		cctx:       ctx,
+		cancel:     cancel,
 		codec:      &json.EventCodec{},
 	}
 
@@ -115,8 +121,7 @@ func (b *EventBus) AddHandler(ctx context.Context, m eh.EventMatcher, h eh.Event
 	b.registered[h.HandlerType()] = struct{}{}
 
 	// Handle until context is cancelled.
-	b.wg.Add(1)
-	go b.handle(ctx, m, h, ch)
+	go b.handle(m, h, ch)
 
 	return nil
 }
@@ -126,10 +131,13 @@ func (b *EventBus) Errors() <-chan eh.EventBusError {
 	return b.errCh
 }
 
-// Wait for all channels to close in the event bus group
-func (b *EventBus) Wait() {
+// Close implements the Close method of the eventhorizon.EventBus interface.
+func (b *EventBus) Close() error {
+	b.cancel()
 	b.wg.Wait()
 	b.group.close()
+
+	return nil
 }
 
 type evt struct {
@@ -138,7 +146,8 @@ type evt struct {
 }
 
 // Handles all events coming in on the channel.
-func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHandler, ch <-chan []byte) {
+func (b *EventBus) handle(m eh.EventMatcher, h eh.EventHandler, ch <-chan []byte) {
+	b.wg.Add(1)
 	defer b.wg.Done()
 
 	for {
@@ -147,7 +156,7 @@ func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHand
 			// Artificial delay to simulate network.
 			time.Sleep(time.Millisecond)
 
-			event, ctx, err := b.codec.UnmarshalEvent(ctx, data)
+			event, ctx, err := b.codec.UnmarshalEvent(b.cctx, data)
 			if err != nil {
 				err = fmt.Errorf("could not unmarshal event: %w", err)
 				select {
@@ -172,7 +181,7 @@ func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHand
 					log.Printf("eventhorizon: missed error in local event bus: %s", err)
 				}
 			}
-		case <-ctx.Done():
+		case <-b.cctx.Done():
 			return
 		}
 	}
diff --git a/eventbus/nats/eventbus.go b/eventbus/nats/eventbus.go
index 8fab0df..8a54701 100644
--- a/eventbus/nats/eventbus.go
+++ b/eventbus/nats/eventbus.go
@@ -39,17 +39,23 @@ type EventBus struct {
 	registered   map[eh.EventHandlerType]struct{}
 	registeredMu sync.RWMutex
 	errCh        chan eh.EventBusError
+	cctx         context.Context
+	cancel       context.CancelFunc
 	wg           sync.WaitGroup
 	codec        eh.EventCodec
 }
 
 // NewEventBus creates an EventBus, with optional settings.
 func NewEventBus(url, appID string, options ...Option) (*EventBus, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+
 	b := &EventBus{
 		appID:      appID,
 		streamName: appID + "_events",
 		registered: map[eh.EventHandlerType]struct{}{},
 		errCh:      make(chan eh.EventBusError, 100),
+		cctx:       ctx,
+		cancel:     cancel,
 		codec:      &json.EventCodec{},
 	}
 
@@ -151,7 +157,7 @@ func (b *EventBus) AddHandler(ctx context.Context, m eh.EventMatcher, h eh.Event
 	// Create a consumer.
 	subject := createConsumerSubject(b.streamName, m)
 	consumerName := fmt.Sprintf("%s_%s", b.appID, h.HandlerType())
-	sub, err := b.js.QueueSubscribe(subject, consumerName, b.handler(ctx, m, h),
+	sub, err := b.js.QueueSubscribe(subject, consumerName, b.handler(b.cctx, m, h),
 		nats.Durable(consumerName),
 		nats.DeliverNew(),
 		nats.ManualAck(),
@@ -167,8 +173,7 @@ func (b *EventBus) AddHandler(ctx context.Context, m eh.EventMatcher, h eh.Event
 	b.registered[h.HandlerType()] = struct{}{}
 
 	// Handle until context is cancelled.
-	b.wg.Add(1)
-	go b.handle(ctx, sub)
+	go b.handle(sub)
 
 	return nil
 }
@@ -178,21 +183,26 @@ func (b *EventBus) Errors() <-chan eh.EventBusError {
 	return b.errCh
 }
 
-// Wait for all channels to close in the event bus group
-func (b *EventBus) Wait() {
+// Close implements the Close method of the eventhorizon.EventBus interface.
+func (b *EventBus) Close() error {
+	b.cancel()
 	b.wg.Wait()
+
 	b.conn.Close()
+
+	return nil
 }
 
 // Handles all events coming in on the channel.
-func (b *EventBus) handle(ctx context.Context, sub *nats.Subscription) {
+func (b *EventBus) handle(sub *nats.Subscription) {
+	b.wg.Add(1)
 	defer b.wg.Done()
 
 	for {
 		select {
-		case <-ctx.Done():
-			if ctx.Err() != context.Canceled {
-				log.Printf("eventhorizon: context error in NATS event bus: %s", ctx.Err())
+		case <-b.cctx.Done():
+			if b.cctx.Err() != context.Canceled {
+				log.Printf("eventhorizon: context error in NATS event bus: %s", b.cctx.Err())
 			}
 			return
 		}
diff --git a/eventbus/redis/eventbus.go b/eventbus/redis/eventbus.go
index 503ce52..e6b0823 100644
--- a/eventbus/redis/eventbus.go
+++ b/eventbus/redis/eventbus.go
@@ -16,6 +16,7 @@ package redis
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"strings"
@@ -39,18 +40,24 @@ type EventBus struct {
 	registered   map[eh.EventHandlerType]struct{}
 	registeredMu sync.RWMutex
 	errCh        chan eh.EventBusError
+	cctx         context.Context
+	cancel       context.CancelFunc
 	wg           sync.WaitGroup
 	codec        eh.EventCodec
 }
 
 // NewEventBus creates an EventBus, with optional settings.
 func NewEventBus(addr, appID, clientID string, options ...Option) (*EventBus, error) {
+	ctx, cancel := context.WithCancel(context.Background())
+
 	b := &EventBus{
 		appID:      appID,
 		clientID:   clientID,
 		streamName: appID + "_events",
 		registered: map[eh.EventHandlerType]struct{}{},
 		errCh:      make(chan eh.EventBusError, 100),
+		cctx:       ctx,
+		cancel:     cancel,
 		codec:      &json.EventCodec{},
 	}
 
@@ -73,8 +80,7 @@ func NewEventBus(addr, appID, clientID string, options ...Option) (*EventBus, er
 
 	// Create client and check connection.
 	b.client = redis.NewClient(b.clientOpts)
-	ctx := context.Background()
-	if res, err := b.client.Ping(ctx).Result(); err != nil || res != "PONG" {
+	if res, err := b.client.Ping(b.cctx).Result(); err != nil || res != "PONG" {
 		return nil, fmt.Errorf("could not check Redis server: %w", err)
 	}
 
@@ -166,8 +172,7 @@ func (b *EventBus) AddHandler(ctx context.Context, m eh.EventMatcher, h eh.Event
 	b.registered[h.HandlerType()] = struct{}{}
 
 	// Handle until context is cancelled.
-	b.wg.Add(1)
-	go b.handle(ctx, m, h, groupName)
+	go b.handle(m, h, groupName)
 
 	return nil
 }
@@ -177,37 +182,39 @@ func (b *EventBus) Errors() <-chan eh.EventBusError {
 	return b.errCh
 }
 
-// Wait for all channels to close in the event bus group
-func (b *EventBus) Wait() {
+// Close implements the Close method of the eventhorizon.EventBus interface.
+func (b *EventBus) Close() error {
+	b.cancel()
 	b.wg.Wait()
-	if err := b.client.Close(); err != nil {
-		log.Printf("eventhorizon: failed to close Redis client: %s", err)
-	}
+
+	return b.client.Close()
 }
 
 // Handles all events coming in on the channel.
-func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHandler, groupName string) {
+func (b *EventBus) handle(m eh.EventMatcher, h eh.EventHandler, groupName string) {
+	b.wg.Add(1)
 	defer b.wg.Done()
 
-	msgHandler := b.handler(m, h, groupName)
+	handler := b.handler(m, h, groupName)
+
 	for {
-		streams, err := b.client.XReadGroup(ctx, &redis.XReadGroupArgs{
+		streams, err := b.client.XReadGroup(b.cctx, &redis.XReadGroupArgs{
 			Group:    groupName,
 			Consumer: groupName + "_" + b.clientID,
 			Streams:  []string{b.streamName, ">"},
 		}).Result()
-		if err != nil && err != context.Canceled {
+		if errors.Is(err, context.Canceled) {
+			break
+		} else if err != nil {
 			err = fmt.Errorf("could not receive: %w", err)
 			select {
-			case b.errCh <- eh.EventBusError{Err: err, Ctx: ctx}:
+			case b.errCh <- eh.EventBusError{Err: err}:
 			default:
 				log.Printf("eventhorizon: missed error in Redis event bus: %s", err)
 			}
 			// Retry the receive loop if there was an error.
 			time.Sleep(time.Second)
 			continue
-		} else if err == context.Canceled {
-			return
 		}
 
 		// Handle all messages from group read.
@@ -216,7 +223,7 @@ func (b *EventBus) handle(ctx context.Context, m eh.EventMatcher, h eh.EventHand
 				continue
 			}
 			for _, msg := range stream.Messages {
-				msgHandler(ctx, &msg)
+				handler(b.cctx, &msg)
 			}
 		}
 	}
diff --git a/eventstore.go b/eventstore.go
index ce2b282..a0da76b 100644
--- a/eventstore.go
+++ b/eventstore.go
@@ -28,6 +28,9 @@ type EventStore interface {
 
 	// Load loads all events for the aggregate id from the store.
 	Load(context.Context, uuid.UUID) ([]Event, error)
+
+	// Close closes the EventStore.
+	Close() error
 }
 
 // EventStoreError is an error in the event store.
diff --git a/eventstore/memory/eventstore.go b/eventstore/memory/eventstore.go
index 425e928..8514753 100644
--- a/eventstore/memory/eventstore.go
+++ b/eventstore/memory/eventstore.go
@@ -178,6 +178,11 @@ type aggregateRecord struct {
 	// Snapshot    eh.Aggregate
 }
 
+// Close implements the Close method of the eventhorizon.EventStore interface.
+func (s *EventStore) Close() error {
+	return nil
+}
+
 // copyEvent duplicates an event.
 func copyEvent(ctx context.Context, event eh.Event) (eh.Event, error) {
 	// Copy data if there is any.
diff --git a/eventstore/mongodb/eventstore.go b/eventstore/mongodb/eventstore.go
index 7323036..34f834a 100644
--- a/eventstore/mongodb/eventstore.go
+++ b/eventstore/mongodb/eventstore.go
@@ -225,12 +225,9 @@ func (s *EventStore) Load(ctx context.Context, id uuid.UUID) ([]eh.Event, error)
 	return events, nil
 }
 
-// Close closes the database client.
-func (s *EventStore) Close(ctx context.Context) error {
-	if err := s.client.Disconnect(ctx); err != nil {
-		return fmt.Errorf("could not close DB connection: %w", err)
-	}
-	return nil
+// Close implements the Close method of the eventhorizon.EventStore interface.
+func (s *EventStore) Close() error {
+	return s.client.Disconnect(context.Background())
 }
 
 // aggregateRecord is the Database representation of an aggregate.
diff --git a/eventstore/mongodb_v2/eventstore.go b/eventstore/mongodb_v2/eventstore.go
index 36017ac..0555057 100644
--- a/eventstore/mongodb_v2/eventstore.go
+++ b/eventstore/mongodb_v2/eventstore.go
@@ -312,12 +312,9 @@ func (s *EventStore) Load(ctx context.Context, id uuid.UUID) ([]eh.Event, error)
 	return events, nil
 }
 
-// Close closes the database client.
-func (s *EventStore) Close(ctx context.Context) error {
-	if err := s.client.Disconnect(ctx); err != nil {
-		return fmt.Errorf("could not close DB connection: %w", err)
-	}
-	return nil
+// Close implements the Close method of the eventhorizon.EventStore interface.
+func (s *EventStore) Close() error {
+	return s.client.Disconnect(context.Background())
 }
 
 // stream is a stream of events, often containing the events for an aggregate.
diff --git a/examples/todomvc/backend/domains/todo/setup.go b/examples/todomvc/backend/domains/todo/setup.go
index 3d631e9..64c2352 100644
--- a/examples/todomvc/backend/domains/todo/setup.go
+++ b/examples/todomvc/backend/domains/todo/setup.go
@@ -15,12 +15,12 @@ import (
 
 // SetupDomain sets up the Todo domain.
 func SetupDomain(
-	ctx context.Context,
 	commandBus *bus.CommandHandler,
 	eventStore eh.EventStore,
 	eventBus eh.EventBus,
 	repo eh.ReadWriteRepo,
 ) error {
+	ctx := context.Background()
 
 	// Set the entity factory for the base repo.
 	if memoryRepo := memory.IntoRepo(ctx, repo); memoryRepo != nil {
diff --git a/examples/todomvc/backend/handler/handler.go b/examples/todomvc/backend/handler/handler.go
index dacc297..572e117 100644
--- a/examples/todomvc/backend/handler/handler.go
+++ b/examples/todomvc/backend/handler/handler.go
@@ -29,7 +29,6 @@ import (
 // NewHandler returns a http.Handler that interacts exposes the command handler,
 // read repo and event bus to the frontend.
 func NewHandler(
-	ctx context.Context,
 	commandHandler eh.CommandHandler,
 	eventBus eh.EventBus,
 	todoRepo eh.ReadRepo,
@@ -40,7 +39,7 @@ func NewHandler(
 	// Add the event bus as a websocket that sends the events as JSON.
 	eventBusHandler := httputils.NewEventBusHandler()
 	observerMiddleware := observer.NewMiddleware(observer.NamedGroup("todomvc"))
-	eventBus.AddHandler(ctx, eh.MatchAll{},
+	eventBus.AddHandler(context.Background(), eh.MatchAll{},
 		eh.UseEventHandlerMiddleware(eventBusHandler, observerMiddleware))
 	h.Handle("/api/events/", eventBusHandler)
 
diff --git a/examples/todomvc/backend/main.go b/examples/todomvc/backend/main.go
index 0d19efc..36fa61f 100644
--- a/examples/todomvc/backend/main.go
+++ b/examples/todomvc/backend/main.go
@@ -91,7 +91,7 @@ func main() {
 	}
 	eventStore = tracingEventStore.NewEventStore(eventStore)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx := context.Background()
 
 	// Add an event logger as an observer.
 	eventLogger := &EventLogger{}
@@ -112,7 +112,7 @@ func main() {
 	todoRepo = tracingRepo.NewRepo(todoRepo)
 
 	// Setup the Todo domain.
-	if err := todo.SetupDomain(ctx, commandBus, eventStore, eventBus, todoRepo); err != nil {
+	if err := todo.SetupDomain(commandBus, eventStore, eventBus, todoRepo); err != nil {
 		log.Fatal("could not setup Todo domain: ", err)
 	}
 
@@ -123,7 +123,7 @@ func main() {
 	)
 
 	// Setup the HTTP handler for commands, read repo and events.
-	h, err := handler.NewHandler(ctx, commandHandler, eventBus, todoRepo, "frontend")
+	h, err := handler.NewHandler(commandHandler, eventBus, todoRepo, "frontend")
 	if err != nil {
 		log.Fatal("could not create handler: ", err)
 	}
@@ -213,9 +213,16 @@ func main() {
 	<-srvClosed
 
 	// Cancel all handlers and wait.
-	cancel()
 	log.Println("waiting for handlers to finish")
-	eventBus.Wait()
+	if err := todoRepo.Close(); err != nil {
+		log.Print("could not close todo repo: ", err)
+	}
+	if err := eventStore.Close(); err != nil {
+		log.Print("could not close event store: ", err)
+	}
+	if err := eventBus.Close(); err != nil {
+		log.Print("could not close event bus: ", err)
+	}
 
 	if err := traceCloser.Close(); err != nil {
 		log.Print("could not close tracer: ", err)
diff --git a/mocks/mocks.go b/mocks/mocks.go
index 49a1d99..58b5f7e 100644
--- a/mocks/mocks.go
+++ b/mocks/mocks.go
@@ -328,7 +328,10 @@ func (m *EventStore) Replace(ctx context.Context, event eh.Event) error {
 	return nil
 }
 
-var _ = eh.EventBus(&EventBus{})
+// Close implements the Close method of the eventhorizon.EventStore interface.
+func (m *EventStore) Close() error {
+	return nil
+}
 
 // EventBus is a mocked eventhorizon.EventBus, useful in testing.
 type EventBus struct {
@@ -338,6 +341,8 @@ type EventBus struct {
 	Err error
 }
 
+var _ = eh.EventBus(&EventBus{})
+
 // HandlerType implements the HandlerType method of the eventhorizon.EventHandler interface.
 func (b *EventBus) HandlerType() eh.EventHandlerType {
 	return "event-bus"
@@ -363,8 +368,10 @@ func (b *EventBus) Errors() <-chan eh.EventBusError {
 	return make(chan eh.EventBusError)
 }
 
-// Wait implements the Wait method of the eventhorizon.EventBus interface.
-func (b *EventBus) Wait() {}
+// Close implements the Close method of the eventhorizon.EventBus interface.
+func (b *EventBus) Close() error {
+	return nil
+}
 
 // Repo is a mocked eventhorizon.ReadRepo, useful in testing.
 type Repo struct {
@@ -436,6 +443,11 @@ func (r *Repo) Remove(ctx context.Context, id uuid.UUID) error {
 	return nil
 }
 
+// Close implements the Close method of the eventhorizon.ReadRepo interface.
+func (r *Repo) Close() error {
+	return nil
+}
+
 type contextKey int
 
 const (
diff --git a/namespace/eventstore.go b/namespace/eventstore.go
index cf032d8..a325acc 100644
--- a/namespace/eventstore.go
+++ b/namespace/eventstore.go
@@ -17,6 +17,7 @@ package namespace
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 
 	eh "github.com/looplab/eventhorizon"
@@ -78,6 +79,25 @@ func (s *EventStore) Load(ctx context.Context, id uuid.UUID) ([]eh.Event, error)
 	return store.Load(ctx, id)
 }
 
+// Close implements the Close method of the eventhorizon.EventStore interface.
+func (s *EventStore) Close() error {
+	s.eventStoresMu.RLock()
+	defer s.eventStoresMu.RUnlock()
+
+	var errStrs []string
+	for _, store := range s.eventStores {
+		if err := store.Close(); err != nil {
+			errStrs = append(errStrs, err.Error())
+		}
+	}
+
+	if len(errStrs) > 0 {
+		return fmt.Errorf("multiple errors: %s", strings.Join(errStrs, ", "))
+	}
+
+	return nil
+}
+
 // eventStore is a helper that returns or creates event stores for each namespace.
 func (s *EventStore) eventStore(ctx context.Context) (eh.EventStore, error) {
 	ns := FromContext(ctx)
diff --git a/namespace/repo.go b/namespace/repo.go
index e0fab68..fcb3b59 100644
--- a/namespace/repo.go
+++ b/namespace/repo.go
@@ -17,6 +17,7 @@ package namespace
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync"
 
 	eh "github.com/looplab/eventhorizon"
@@ -124,6 +125,25 @@ func (r *Repo) Remove(ctx context.Context, id uuid.UUID) error {
 	return repo.Remove(ctx, id)
 }
 
+// Close implements the Close method of the eventhorizon.WriteRepo interface.
+func (r *Repo) Close() error {
+	r.reposMu.RLock()
+	defer r.reposMu.RUnlock()
+
+	var errStrs []string
+	for _, repo := range r.repos {
+		if err := repo.Close(); err != nil {
+			errStrs = append(errStrs, err.Error())
+		}
+	}
+
+	if len(errStrs) > 0 {
+		return fmt.Errorf("multiple errors: %s", strings.Join(errStrs, ", "))
+	}
+
+	return nil
+}
+
 // repo is a helper that returns or creates repos for each namespace.
 func (r *Repo) repo(ctx context.Context) (eh.ReadWriteRepo, error) {
 	ns := FromContext(ctx)
diff --git a/repo.go b/repo.go
index 21aa4fa..0297b5f 100644
--- a/repo.go
+++ b/repo.go
@@ -32,6 +32,9 @@ type ReadRepo interface {
 
 	// FindAll returns all entities in the repository.
 	FindAll(context.Context) ([]Entity, error)
+
+	// Close closes the ReadRepo.
+	Close() error
 }
 
 // WriteRepo is a write repository for entities.
diff --git a/repo/memory/repo.go b/repo/memory/repo.go
index 7a7644d..26d0727 100644
--- a/repo/memory/repo.go
+++ b/repo/memory/repo.go
@@ -185,3 +185,8 @@ func (r *Repo) Remove(ctx context.Context, id uuid.UUID) error {
 func (r *Repo) SetEntityFactory(f func() eh.Entity) {
 	r.factoryFn = f
 }
+
+// Close implements the Close method of the eventhorizon.WriteRepo interface.
+func (r *Repo) Close() error {
+	return nil
+}
diff --git a/repo/mongodb/repo.go b/repo/mongodb/repo.go
index 46da1aa..bc4219b 100644
--- a/repo/mongodb/repo.go
+++ b/repo/mongodb/repo.go
@@ -352,7 +352,7 @@ func (r *Repo) Clear(ctx context.Context) error {
 	return nil
 }
 
-// Close closes a database session.
-func (r *Repo) Close(ctx context.Context) {
-	r.client.Disconnect(ctx)
+// Close implements the Close method of the eventhorizon.WriteRepo interface.
+func (r *Repo) Close() error {
+	return r.client.Disconnect(context.Background())
 }


Test patch:
diff --git a/eventbus/acceptance_testing.go b/eventbus/acceptance_testing.go
index e2309fe..16b6ccf 100644
--- a/eventbus/acceptance_testing.go
+++ b/eventbus/acceptance_testing.go
@@ -64,7 +64,7 @@ func TestAddHandler(t *testing.T, bus1 eh.EventBus) {
 //   }
 //
 func AcceptanceTest(t *testing.T, bus1, bus2 eh.EventBus, timeout time.Duration) {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx := context.Background()
 	ctx = mocks.WithContextOne(ctx, "testval")
 
 	// Without handler.
@@ -245,10 +245,14 @@ func AcceptanceTest(t *testing.T, bus1, bus2 eh.EventBus, timeout time.Duration)
 		}
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	bus1.Wait()
-	bus2.Wait()
+	// TODO: Test context cancellation.
+
+	if err := bus1.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
+	if err := bus2.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func checkBusErrors(t *testing.T, bus eh.EventBus) {
@@ -357,8 +361,9 @@ func benchmark(t bench, bus eh.EventBus, numAggregates, numHandlers, numEvents i
 
 	wg.Wait() // Wait for all events to be received.
 	t.StopTimer()
+	cancel() // Stop receiving goroutines.
 
-	// Cancel handler and wait.
-	cancel()
-	bus.Wait()
+	if err := bus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
diff --git a/eventstore/mongodb/eventmaintenance_test.go b/eventstore/mongodb/eventmaintenance_test.go
index 2b0b14f..54c08f6 100644
--- a/eventstore/mongodb/eventmaintenance_test.go
+++ b/eventstore/mongodb/eventmaintenance_test.go
@@ -51,7 +51,7 @@ func TestEventStoreMaintenanceIntegration(t *testing.T) {
 	if store == nil {
 		t.Fatal("there should be a store")
 	}
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	eventstore.MaintenanceAcceptanceTest(t, store, store, context.Background())
 }
diff --git a/eventstore/mongodb/eventstore_test.go b/eventstore/mongodb/eventstore_test.go
index 3528809..cd482e9 100644
--- a/eventstore/mongodb/eventstore_test.go
+++ b/eventstore/mongodb/eventstore_test.go
@@ -55,7 +55,7 @@ func TestEventStoreIntegration(t *testing.T) {
 	if store == nil {
 		t.Fatal("there should be a store")
 	}
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	eventstore.AcceptanceTest(t, store, context.Background())
 }
@@ -91,7 +91,7 @@ func TestWithEventHandlerIntegration(t *testing.T) {
 	if store == nil {
 		t.Fatal("there should be a store")
 	}
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	ctx := context.Background()
 
@@ -159,7 +159,7 @@ func BenchmarkEventStore(b *testing.B) {
 	if store == nil {
 		b.Fatal("there should be a store")
 	}
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	eventstore.Benchmark(b, store)
 }
diff --git a/eventstore/mongodb_v2/eventmaintenance_test.go b/eventstore/mongodb_v2/eventmaintenance_test.go
index 411efbe..0525b87 100644
--- a/eventstore/mongodb_v2/eventmaintenance_test.go
+++ b/eventstore/mongodb_v2/eventmaintenance_test.go
@@ -51,7 +51,7 @@ func TestEventStoreMaintenanceIntegration(t *testing.T) {
 	if store == nil {
 		t.Fatal("there should be a store")
 	}
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	eventstore.MaintenanceAcceptanceTest(t, store, store, context.Background())
 }
diff --git a/eventstore/mongodb_v2/eventstore_test.go b/eventstore/mongodb_v2/eventstore_test.go
index abd6c29..97f1048 100644
--- a/eventstore/mongodb_v2/eventstore_test.go
+++ b/eventstore/mongodb_v2/eventstore_test.go
@@ -56,7 +56,7 @@ func TestEventStoreIntegration(t *testing.T) {
 		t.Fatal("there should be a store")
 	}
 
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	eventstore.AcceptanceTest(t, store, context.Background())
 }
@@ -92,7 +92,7 @@ func TestWithEventHandlerIntegration(t *testing.T) {
 	if store == nil {
 		t.Fatal("there should be a store")
 	}
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	ctx := context.Background()
 
@@ -164,7 +164,7 @@ func BenchmarkEventStore(b *testing.B) {
 		b.Fatal("there should be a store")
 	}
 
-	defer store.Close(context.Background())
+	defer store.Close()
 
 	eventstore.Benchmark(b, store)
 }
diff --git a/examples/guestlist/memory/memory_test.go b/examples/guestlist/memory/memory_test.go
index 0182608..ab07ee5 100644
--- a/examples/guestlist/memory/memory_test.go
+++ b/examples/guestlist/memory/memory_test.go
@@ -59,8 +59,7 @@ func Example() {
 	guestListRepo := memory.NewRepo()
 	guestListRepo.SetEntityFactory(func() eh.Entity { return &guestlist.GuestList{} })
 
-	// Create a context with cancellation.
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx := context.Background()
 
 	// Setup a test utility waiter that waits for all 11 events to occur before
 	// evaluating results.
@@ -164,9 +163,18 @@ func Example() {
 			l.NumGuests, l.NumAccepted, l.NumDeclined, l.NumConfirmed, l.NumDenied)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := invitationRepo.Close(); err != nil {
+		log.Println("error closing invitation repo:", err)
+	}
+	if err := guestListRepo.Close(); err != nil {
+		log.Println("error closing guest list repo:", err)
+	}
+	if err := eventStore.Close(); err != nil {
+		log.Println("error closing event store:", err)
+	}
+	if err := eventBus.Close(); err != nil {
+		log.Println("error closing event bus:", err)
+	}
 
 	// Output:
 	// invitation: Athena - confirmed
diff --git a/examples/guestlist/mongodb/mongodb_test.go b/examples/guestlist/mongodb/mongodb_test.go
index 59473f5..19b93d4 100644
--- a/examples/guestlist/mongodb/mongodb_test.go
+++ b/examples/guestlist/mongodb/mongodb_test.go
@@ -88,8 +88,7 @@ guest list: 4 invited - 3 accepted, 1 declined - 2 confirmed, 1 denied`)
 	}
 	guestListRepo.SetEntityFactory(func() eh.Entity { return &guestlist.GuestList{} })
 
-	// Create a context with cancellation.
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx := context.Background()
 
 	// Setup a test utility waiter that waits for all 11 events to occur before
 	// evaluating results.
@@ -198,9 +197,18 @@ guest list: 4 invited - 3 accepted, 1 declined - 2 confirmed, 1 denied`)
 			l.NumGuests, l.NumAccepted, l.NumDeclined, l.NumConfirmed, l.NumDenied)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := invitationRepo.Close(); err != nil {
+		log.Println("error closing invitation repo:", err)
+	}
+	if err := guestListRepo.Close(); err != nil {
+		log.Println("error closing guest list repo:", err)
+	}
+	if err := eventStore.Close(); err != nil {
+		log.Println("error closing event store:", err)
+	}
+	if err := eventBus.Close(); err != nil {
+		log.Println("error closing event bus:", err)
+	}
 
 	// Output:
 	// invitation: Athena - confirmed
diff --git a/examples/todomvc/backend/handler/handler_test.go b/examples/todomvc/backend/handler/handler_test.go
index 83acd78..a515020 100644
--- a/examples/todomvc/backend/handler/handler_test.go
+++ b/examples/todomvc/backend/handler/handler_test.go
@@ -43,11 +43,9 @@ import (
 )
 
 func TestStaticFiles(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
-
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -59,9 +57,9 @@ func TestStaticFiles(t *testing.T) {
 		t.Error("there should be a 200 status for /")
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestGetAll(t *testing.T) {
@@ -69,11 +67,9 @@ func TestGetAll(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
-
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -88,6 +84,8 @@ func TestGetAll(t *testing.T) {
 		t.Error("the body should be correct:", w.Body.String())
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -119,9 +117,9 @@ func TestGetAll(t *testing.T) {
 		t.Error("the body should be correct:", w.Body.String())
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestCreate(t *testing.T) {
@@ -129,11 +127,9 @@ func TestCreate(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -150,6 +146,8 @@ func TestCreate(t *testing.T) {
 		t.Error("the body should be correct:", w.Body.String())
 	}
 
+	ctx := context.Background()
+
 	waiter := waiter.NewEventHandler()
 	eventBus.AddHandler(ctx, eh.MatchEvents{todo.Created},
 		eh.UseEventHandlerMiddleware(waiter, observer.Middleware))
@@ -179,9 +177,9 @@ func TestCreate(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestDelete(t *testing.T) {
@@ -189,15 +187,15 @@ func TestDelete(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -231,9 +229,9 @@ func TestDelete(t *testing.T) {
 		t.Error("there should be a not found error:", err)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestAddItem(t *testing.T) {
@@ -241,15 +239,15 @@ func TestAddItem(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
-
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -302,9 +300,9 @@ func TestAddItem(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestRemoveItem(t *testing.T) {
@@ -312,15 +310,15 @@ func TestRemoveItem(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -374,9 +372,9 @@ func TestRemoveItem(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestRemoveCompleted(t *testing.T) {
@@ -384,15 +382,15 @@ func TestRemoveCompleted(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
-
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -466,9 +464,9 @@ func TestRemoveCompleted(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestSetItemDesc(t *testing.T) {
@@ -476,15 +474,15 @@ func TestSetItemDesc(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -543,9 +541,9 @@ func TestSetItemDesc(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestCheckItem(t *testing.T) {
@@ -553,15 +551,15 @@ func TestCheckItem(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -631,9 +629,9 @@ func TestCheckItem(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
 func TestCheckAllItems(t *testing.T) {
@@ -641,15 +639,15 @@ func TestCheckAllItems(t *testing.T) {
 		return time.Date(2017, time.July, 10, 23, 0, 0, 0, time.UTC)
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	commandHandler, eventBus, todoRepo := NewTestSession()
 
-	commandHandler, eventBus, todoRepo := NewTestSession(ctx)
-
-	h, err := NewHandler(ctx, commandHandler, eventBus, todoRepo, "../../frontend")
+	h, err := NewHandler(commandHandler, eventBus, todoRepo, "../../frontend")
 	if err != nil {
 		t.Fatal(err)
 	}
 
+	ctx := context.Background()
+
 	id := uuid.New()
 	if err := commandHandler.HandleCommand(ctx, &todo.Create{
 		ID: id,
@@ -722,12 +720,12 @@ func TestCheckAllItems(t *testing.T) {
 		t.Log("expected:", expected)
 	}
 
-	// Cancel all handlers and wait.
-	cancel()
-	eventBus.Wait()
+	if err := eventBus.Close(); err != nil {
+		t.Error("there should be no error:", err)
+	}
 }
 
-func NewTestSession(ctx context.Context) (
+func NewTestSession() (
 	eh.CommandHandler,
 	eh.EventBus,
 	eh.ReadWriteRepo,
@@ -745,7 +743,7 @@ func NewTestSession(ctx context.Context) (
 
 	// Create the read repositories.
 	todoRepo := memory.NewRepo()
-	if err := todo.SetupDomain(ctx, commandBus, eventStore, eventBus, todoRepo); err != nil {
+	if err := todo.SetupDomain(commandBus, eventStore, eventBus, todoRepo); err != nil {
 		log.Println("could not setup domain:", err)
 	}
 
@@ -800,7 +798,7 @@ func NewIntegrationTestSession(ctx context.Context) (
 	}
 
 	// Setup the todo domain.
-	if err := todo.SetupDomain(ctx, commandBus, eventStore, eventBus, todoRepo); err != nil {
+	if err := todo.SetupDomain(commandBus, eventStore, eventBus, todoRepo); err != nil {
 		log.Println("could not setup domain:", err)
 	}
 
diff --git a/repo/mongodb/repo_test.go b/repo/mongodb/repo_test.go
index 9d18dfe..80636dd 100644
--- a/repo/mongodb/repo_test.go
+++ b/repo/mongodb/repo_test.go
@@ -60,7 +60,7 @@ func TestReadRepoIntegration(t *testing.T) {
 	if r == nil {
 		t.Error("there should be a repository")
 	}
-	defer r.Close(context.Background())
+	defer r.Close()
 
 	r.SetEntityFactory(func() eh.Entity {
 		return &mocks.Model{}
@@ -193,7 +193,7 @@ func TestIntoRepo(t *testing.T) {
 	if err != nil {
 		t.Error("there should be no error:", err)
 	}
-	defer inner.Close(context.Background())
+	defer inner.Close()
 
 	outer := &mocks.Repo{ParentRepo: inner}
 	if r := IntoRepo(context.Background(), outer); r != inner {