feat: ensure shutdowns are graceful (#42)

Author: w-h-a

cmd/engine.go

@@ -34,25 +34,26 @@ func StartEngine(ctx *cli.Context) error {
 	numServices := 0
 
 	// worker
+	var workerServer serverv2.Server
 	var w *worker.Service
 	if config.Mode() == "standalone" || config.Mode() == "worker" {
 		runnerClient := initRunner()
 
-		w = engine.NewWorker(
+		workerServer, w = engine.NewWorker(
 			brokerClient,
 			runnerClient,
 		)
 
-		numServices++
+		numServices += 2
 	}
 
 	// coordinator
-	var httpServer serverv2.Server
+	var coordinatorServer serverv2.Server
 	var c *coordinator.Service
 	if config.Mode() == "standalone" || config.Mode() == "coordinator" {
 		readwriterClient := initReadWriter()
 
-		httpServer, c = engine.NewCoordinator(
+		coordinatorServer, c = engine.NewCoordinator(
 			brokerClient,
 			readwriterClient,
 		)
@@ -74,6 +75,13 @@ func StartEngine(ctx *cli.Context) error {
 			slog.InfoContext(ctx.Context, "starting worker")
 			ch <- w.Start(stop)
 		}()
+
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			slog.InfoContext(ctx.Context, "starting worker http server", "address", config.WorkerHttp())
+			ch <- workerServer.Start()
+		}()
 	}
 
 	// start coordinator
@@ -91,8 +99,8 @@ func StartEngine(ctx *cli.Context) error {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			slog.InfoContext(ctx.Context, "starting http server", "address", config.HttpAddress())
-			ch <- httpServer.Start()
+			slog.InfoContext(ctx.Context, "starting coordinator http server", "address", config.CoordinatorHttp())
+			ch <- coordinatorServer.Start()
 		}()
 	}
 

docker-compose.yaml

@@ -31,11 +31,14 @@ services:
     container_name: worker
     build: .
     command: engine
+    ports:
+      - "4001:4001"
     environment:
       - ENV=dev
       - NAME=worker
       - VERSION=0.1.0-alpha.0
       - MODE=worker
+      - HTTP_ADDRESS=:4001
       - BROKER=rabbit
       - BROKER_LOCATION=amqp://guest:guest@rabbitmq:5672
       - RUNNER=docker
@@ -56,8 +59,8 @@ services:
       - ENV=dev
       - NAME=coordinator
       - VERSION=0.1.0-alpha.0
-      - HTTP_ADDRESS=:4000
       - MODE=coordinator
+      - HTTP_ADDRESS=:4000
       - BROKER=rabbit
       - BROKER_LOCATION=amqp://guest:guest@rabbitmq:5672
       - READ_WRITER=postgres

internal/engine/clients/broker/broker.go

@@ -17,6 +17,7 @@ var (
 )
 
 type Broker interface {
-	Subscribe(ctx context.Context, callback func(ctx context.Context, data []byte) error, opts ...SubscribeOption) (chan struct{}, error)
+	Subscribe(ctx context.Context, callback func(ctx context.Context, data []byte) error, opts ...SubscribeOption) error
 	Publish(ctx context.Context, data []byte, opts ...PublishOption) error
+	Close(ctx context.Context) error
 }

internal/engine/clients/broker/memory/broker.go

@@ -12,9 +12,12 @@ type memoryBroker struct {
 	options broker.Options
 	queues  map[string]chan []byte
 	mtx     sync.RWMutex
+	exit    chan struct{}
+	wg      sync.WaitGroup
+	once    sync.Once
 }
 
-func (b *memoryBroker) Subscribe(ctx context.Context, callback func(ctx context.Context, data []byte) error, opts ...broker.SubscribeOption) (chan struct{}, error) {
+func (b *memoryBroker) Subscribe(ctx context.Context, callback func(ctx context.Context, data []byte) error, opts ...broker.SubscribeOption) error {
 	options := broker.NewSubscribeOptions(opts...)
 
 	// span
@@ -28,12 +31,13 @@ func (b *memoryBroker) Subscribe(ctx context.Context, callback func(ctx context.
 	}
 	b.mtx.Unlock()
 
-	ch := make(chan struct{})
-
+	b.wg.Add(1)
 	go func() {
+		defer b.wg.Done()
+
 		for {
 			select {
-			case <-ch:
+			case <-b.exit:
 				return
 			case data := <-q:
 				if err := callback(ctx, data); err != nil {
@@ -44,7 +48,7 @@ func (b *memoryBroker) Subscribe(ctx context.Context, callback func(ctx context.
 		}
 	}()
 
-	return ch, nil
+	return nil
 }
 
 func (b *memoryBroker) Publish(ctx context.Context, data []byte, opts ...broker.PublishOption) error {
@@ -66,13 +70,36 @@ func (b *memoryBroker) Publish(ctx context.Context, data []byte, opts ...broker.
 	return nil
 }
 
+func (b *memoryBroker) Close(ctx context.Context) error {
+	done := make(chan struct{})
+
+	b.once.Do(func() {
+		close(b.exit)
+
+		go func() {
+			b.wg.Wait()
+			close(done)
+		}()
+	})
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		return nil
+	}
+}
+
 func NewBroker(opts ...broker.Option) broker.Broker {
 	options := broker.NewOptions(opts...)
 
 	b := &memoryBroker{
 		options: options,
 		queues:  map[string]chan []byte{},
 		mtx:     sync.RWMutex{},
+		exit:    make(chan struct{}),
+		wg:      sync.WaitGroup{},
+		once:    sync.Once{},
 	}
 
 	return b