Account for closed connection when publishing messages

Author: exanubes

client.go

@@ -3,6 +3,7 @@ package appsync
 import (
 	"context"
 
+	"github.com/exanubes/appsync/internal/app/lifecycle"
 	"github.com/exanubes/appsync/internal/app/usecases/publish"
 	"github.com/exanubes/appsync/internal/app/usecases/subscribe"
 	"github.com/exanubes/appsync/internal/composition"
@@ -13,10 +14,15 @@ import (
 const ProtocolEvents = "aws-appsync-event-ws"
 
 type appsync_client struct {
-	usecases *composition.UseCases
+	usecases   *composition.UseCases
+	connection *lifecycle.State
 }
 
 func (client *appsync_client) Publish(ctx context.Context, input PublishCommandInput) error {
+	if err := client.connection.Err(); err != nil {
+		return err
+	}
+
 	frame := &events.FrameBuilder{}
 
 	err := client.usecases.Publish.Publish(ctx, publish.PublishCommandInput{
@@ -28,6 +34,10 @@ func (client *appsync_client) Publish(ctx context.Context, input PublishCommandI
 }
 
 func (client *appsync_client) Subscribe(ctx context.Context, input SubscribeCommandInput) (Subscription, error) {
+	if err := client.connection.Err(); err != nil {
+		return nil, err
+	}
+
 	frame := &events.FrameBuilder{}
 	result, err := client.usecases.Subscribe.Execute(ctx, subscribe.SubscribeCommandInput{
 		Channel: input.Channel,

connect.go

@@ -138,8 +138,8 @@ func (builder *builder) Connect(ctx context.Context) (*appsync_client, error) {
 	clock := clock.New()
 	heartbeat := heartbeat.New(clock)
 	ingress_queue := queue.NewIngressQueue(builder.backpressure.ConnectionInbound)
-	egress_queue := queue.NewEgressQueue(builder.backpressure.ConnectionOutbound)
-	pending_registry := pending.NewRegistry()
+	egress_queue := queue.NewEgressQueue(builder.backpressure.ConnectionOutbound, connection_state)
+	pending_registry := pending.NewRegistry(connection_state)
 	io_loops := io.New(ingress_queue, egress_queue, connection_output.Connection, msg_codec)
 	usecases, services := composition.NewUseCases(
 		request_authorizer,
@@ -172,7 +172,8 @@ func (builder *builder) Connect(ctx context.Context) (*appsync_client, error) {
 	usecases.Shutdown = shutdown_connection_usecase
 
 	return &appsync_client{
-		usecases: usecases,
+		usecases:   usecases,
+		connection: connection_state,
 	}, nil
 }
 

internal/app/pending/pending.go

@@ -1,14 +1,20 @@
 package pending
 
-import "context"
+import (
+	"context"
+
+	"github.com/exanubes/appsync/internal/app"
+)
 
 type Registry struct {
-	store map[string]chan error
+	store      map[string]chan error
+	connection ConnectionState
 }
 
-func NewRegistry() *Registry {
+func NewRegistry(connection ConnectionState) *Registry {
 	return &Registry{
-		store: make(map[string]chan error),
+		store:      make(map[string]chan error),
+		connection: connection,
 	}
 }
 
@@ -33,9 +39,15 @@ func (registry Registry) Fulfill(ctx context.Context, id string, err error) erro
 		return ctx.Err()
 	case reply <- err:
 		return nil
+	case <-registry.connection.Done():
+		return app.ErrConnectionClosed
 	}
 }
 
+func (registry Registry) Remove(id string) {
+	delete(registry.store, id)
+}
+
 func (registry Registry) get(id string) chan error {
 	return registry.store[id]
 }
@@ -54,5 +66,7 @@ func (registry Registry) Consume(ctx context.Context, id string) error {
 		return ctx.Err()
 	case res := <-reply:
 		return res
+	case <-registry.connection.Done():
+		return app.ErrConnectionClosed
 	}
 }

internal/app/pending/pending_test.go

@@ -8,6 +8,10 @@ import (
 	"github.com/exanubes/appsync/internal/app/pending"
 )
 
+type mock_connection_state struct{}
+
+func (m *mock_connection_state) Done() <-chan struct{} { return make(chan struct{}) }
+
 func TestHas(t *testing.T) {
 	tests := []struct {
 		name       string
@@ -33,7 +37,7 @@ func TestHas(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			registry := pending.NewRegistry()
+			registry := pending.NewRegistry(&mock_connection_state{})
 			tt.setup(registry)
 			has := registry.Has(tt.id)
 			if has != tt.expect_has {
@@ -56,7 +60,7 @@ func TestRegister(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			registry := pending.NewRegistry()
+			registry := pending.NewRegistry(&mock_connection_state{})
 			registry.Register(tt.id)
 			if !registry.Has(tt.id) {
 				t.Errorf("Has(%q) = false after Register, want true", tt.id)
@@ -123,7 +127,7 @@ func TestFulfill(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			registry := pending.NewRegistry()
+			registry := pending.NewRegistry(&mock_connection_state{})
 			tt.setup(registry)
 			err := registry.Fulfill(tt.ctx(), tt.id, tt.err)
 			if !errors.Is(err, tt.expect_err) {
@@ -203,7 +207,7 @@ func TestConsume(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			registry := pending.NewRegistry()
+			registry := pending.NewRegistry(&mock_connection_state{})
 			tt.setup(registry)
 			err := registry.Consume(tt.ctx(), tt.id)
 			if !errors.Is(err, tt.expect_err) {

internal/app/pending/ports.go

@@ -0,0 +1,5 @@
+package pending
+
+type ConnectionState interface {
+	Done() <-chan struct{}
+}

internal/app/queue/egress-queue.go

@@ -2,15 +2,19 @@ package queue
 
 import (
 	"context"
+
+	"github.com/exanubes/appsync/internal/app"
 )
 
 type EgressQueue struct {
-	inbox chan []byte
+	inbox      chan []byte
+	connection ConnectionState
 }
 
-func NewEgressQueue(max_size uint) *EgressQueue {
+func NewEgressQueue(max_size uint, connection ConnectionState) *EgressQueue {
 	return &EgressQueue{
-		inbox: make(chan []byte, max_size),
+		inbox:      make(chan []byte, max_size),
+		connection: connection,
 	}
 }
 
@@ -29,5 +33,7 @@ func (registry *EgressQueue) Enqueue(ctx context.Context, payload []byte) error
 		return nil
 	case <-ctx.Done():
 		return ctx.Err()
+	case <-registry.connection.Done():
+		return app.ErrConnectionClosed
 	}
 }

internal/app/queue/egress_queue_test.go

@@ -8,6 +8,10 @@ import (
 	"github.com/exanubes/appsync/internal/app/queue"
 )
 
+type mock_connection_state struct{}
+
+func (m *mock_connection_state) Done() <-chan struct{} { return make(chan struct{}) }
+
 func TestEgressQueue_Next(t *testing.T) {
 	payload := []byte("test-payload")
 
@@ -42,7 +46,7 @@ func TestEgressQueue_Next(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			q := queue.NewEgressQueue(1)
+			q := queue.NewEgressQueue(1, &mock_connection_state{})
 			tt.setup(q)
 
 			got, err := q.Next(tt.ctx())
@@ -91,7 +95,7 @@ func TestEgressQueue_Enqueue(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			q := queue.NewEgressQueue(1)
+			q := queue.NewEgressQueue(1, &mock_connection_state{})
 			tt.setup(q)
 
 			err := q.Enqueue(tt.ctx(), tt.payload)

internal/app/queue/ports.go

@@ -0,0 +1,5 @@
+package queue
+
+type ConnectionState interface {
+	Done() <-chan struct{}
+}

internal/app/router/router_test.go

@@ -10,6 +10,10 @@ import (
 	"github.com/exanubes/appsync/internal/app/router"
 )
 
+type mock_connection_state struct{}
+
+func (m *mock_connection_state) Done() <-chan struct{} { return make(chan struct{}) }
+
 type mock_receive_data struct {
 	called   bool
 	received protocol.DataMessage
@@ -110,7 +114,7 @@ func TestHandle(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			registry := pending.NewRegistry()
+			registry := pending.NewRegistry(&mock_connection_state{})
 			if tt.setup != nil {
 				tt.setup(registry)
 			}

internal/app/services/request/ports.go

@@ -10,6 +10,7 @@ type Outbox interface {
 
 type Registry interface {
 	Register(id string)
+	Remove(id string)
 	Consume(ctx context.Context, id string) error
 	Has(id string) bool
 }