fix: graceful shutdown of the service (#19)

close the connection during shutdown and capture the data loss metric appropriately

Author: siddhanta-rath

app/proc.go

@@ -12,10 +12,10 @@ func Run() error {
 	ctx, cancel := context.WithCancel(ctx)
 
 	//@TODO - init config
-
+	shutdown := make(chan bool)
 	//start server
-	StartServer(ctx, cancel)
+	StartServer(ctx, cancel, shutdown)
 	logger.Info("App.Run --> Complete")
-	<-ctx.Done()
+	<-shutdown
 	return nil
 }

app/server.go

@@ -19,9 +19,9 @@ import (
 )
 
 // StartServer starts the server
-func StartServer(ctx context.Context, cancel context.CancelFunc) {
+func StartServer(ctx context.Context, cancel context.CancelFunc, shutdown chan bool) {
 	bufferChannel := make(chan collection.CollectRequest, config.Worker.ChannelSize)
-	httpServices := services.Create(bufferChannel)
+	httpServices := services.Create(bufferChannel, ctx)
 	logger.Info("Start Server -->")
 	httpServices.Start(ctx, cancel)
 	logger.Info("Start publisher -->")
@@ -37,38 +37,47 @@ func StartServer(ctx context.Context, cancel context.CancelFunc) {
 	workerPool.StartWorkers()
 	go kPublisher.ReportStats()
 	go reportProcMetrics()
-	go shutDownServer(ctx, cancel, httpServices, bufferChannel, workerPool, kPublisher)
+	// create signal channel at startup
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+
+	go shutDownServer(ctx, cancel, httpServices, bufferChannel, workerPool, kPublisher, shutdown, signalChan)
 }
 
-func shutDownServer(ctx context.Context, cancel context.CancelFunc, httpServices services.Services, bufferChannel chan collection.CollectRequest, workerPool *worker.Pool, kp *publisher.Kafka) {
-	signalChan := make(chan os.Signal)
-	signal.Notify(signalChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+func shutDownServer(ctx context.Context, cancel context.CancelFunc, httpServices services.Services, bufferChannel chan collection.CollectRequest,
+	workerPool *worker.Pool, kp *publisher.Kafka, shutdown chan bool, signalChan chan os.Signal) {
 	for {
 		sig := <-signalChan
 		switch sig {
 		case syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT:
 			logger.Info(fmt.Sprintf("[App.Server] Received a signal %s", sig))
 			httpServices.Shutdown(ctx)
 			logger.Info("Server shutdown all the listeners")
+			cancel()
+			close(bufferChannel)
 			timedOut := workerPool.FlushWithTimeOut(config.Worker.WorkerFlushTimeout)
 			if timedOut {
 				logger.Info(fmt.Sprintf("WorkerPool flush timedout %t", timedOut))
+			} else {
+				logger.Info("WorkerPool flushed all events")
 			}
 			flushInterval := config.PublisherKafka.FlushInterval
 			logger.Info("Closing Kafka producer")
 			logger.Info(fmt.Sprintf("Wait %d ms for all messages to be delivered", flushInterval))
 			eventsInProducer := kp.Close()
-			/**
-			@TODO - should compute the actual no., of events per batch and therefore the total. We can do this only when we close all the active connections
-			Until then we fall back to approximation */
-			eventsInChannel := len(bufferChannel) * 7
-			logger.Info(fmt.Sprintf("Outstanding unprocessed events in the channel, data lost ~ (No batches %d * 5 events) = ~%d", len(bufferChannel), eventsInChannel))
-			metrics.Count("kafka_messages_delivered_total", eventsInChannel+eventsInProducer, "success=false")
+			eventCountInChannel := 0
+			for i := 0; i < len(bufferChannel); i++ {
+				req := <-bufferChannel
+				eventCountInChannel += len(req.Events)
+			}
+			logger.Info(fmt.Sprintf("number of events dropped during the shutdown %d", eventCountInChannel+eventsInProducer))
+			metrics.Count("total_data_loss", eventCountInChannel+eventsInProducer, "reason=shutdown")
 			logger.Info("Exiting server")
-			cancel()
+			shutdown <- true
 		default:
 			logger.Info(fmt.Sprintf("[App.Server] Received a unexpected signal %s", sig))
 		}
+		return
 	}
 }
 

app/server_test.go

@@ -0,0 +1,101 @@
+package app
+
+import (
+	"context"
+	"github.com/goto/raccoon/publisher"
+	"github.com/goto/raccoon/services"
+	"github.com/goto/raccoon/worker"
+	"gopkg.in/confluentinc/confluent-kafka-go.v1/kafka"
+	"os"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/goto/raccoon/collection"
+)
+
+func TestShutDownServer(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	mockKafka := &mockKafkaClient{}
+
+	kp := publisher.NewKafkaFromClient(mockKafka, 50, "test")
+
+	shutdownCh := make(chan bool, 1)
+	bufferCh := make(chan collection.CollectRequest, 1)
+
+	services := services.Create(bufferCh, ctx)
+
+	wp := worker.CreateWorkerPool(1, bufferCh, 1, kp)
+
+	// run shutdown in background
+	sigCh := make(chan os.Signal, 1)
+	go shutDownServer(ctx, cancel, services, bufferCh, wp, kp, shutdownCh, sigCh)
+
+	// send a termination signal after short delay
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		sigCh <- syscall.SIGTERM
+	}()
+
+	select {
+	case <-shutdownCh:
+		t.Log("shutdown executed successfully")
+	case <-time.After(500 * time.Millisecond):
+		t.Error("shutdown execution failed")
+	}
+
+	if !isClosed(bufferCh) {
+		t.Errorf("expected buffer channel to be closed")
+	}
+	if !mockKafka.FlushCalled {
+		t.Errorf("expected Kafka.FlushCalled to be called")
+	}
+	if !mockKafka.CloseCalled {
+		t.Errorf("expected Kafka.CloseCalled to be called")
+	}
+}
+
+func isClosed(ch <-chan collection.CollectRequest) bool {
+	select {
+	case _, ok := <-ch:
+		return !ok
+	default:
+		return false
+	}
+}
+
+// ---- Mocks ----
+// mockKafkaClient is a mock for the Client interface
+type mockKafkaClient struct {
+	// Tracking flags
+	ProduceCalled bool
+	CloseCalled   bool
+	FlushCalled   bool
+	EventsCalled  bool
+
+	ReturnFlushLeft int
+	EventChan       chan kafka.Event
+}
+
+func (m *mockKafkaClient) Produce(msg *kafka.Message, deliveryChan chan kafka.Event) error {
+	m.ProduceCalled = true
+	return nil
+}
+
+func (m *mockKafkaClient) Close() {
+	m.CloseCalled = true
+}
+
+func (m *mockKafkaClient) Flush(timeout int) int {
+	m.FlushCalled = true
+	return m.ReturnFlushLeft
+}
+
+func (m *mockKafkaClient) Events() chan kafka.Event {
+	m.EventsCalled = true
+	if m.EventChan == nil {
+		m.EventChan = make(chan kafka.Event, 1)
+	}
+	return m.EventChan
+}

services/rest/service.go

@@ -19,10 +19,10 @@ type Service struct {
 	s         *http.Server
 }
 
-func NewRestService(c collection.Collector) *Service {
+func NewRestService(c collection.Collector, ctx context.Context) *Service {
 	pingChannel := make(chan connection.Conn, config.ServerWs.ServerMaxConn)
 	wh := websocket.NewHandler(pingChannel, c)
-	go websocket.Pinger(pingChannel, config.ServerWs.PingerSize, config.ServerWs.PingInterval, config.ServerWs.WriteWaitInterval)
+	go websocket.Pinger(ctx, pingChannel, config.ServerWs.PingerSize, config.ServerWs.PingInterval, config.ServerWs.WriteWaitInterval)
 
 	go reportConnectionMetrics(*wh.Table())
 

services/rest/websocket/pinger.go

@@ -1,6 +1,7 @@
 package websocket
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -10,14 +11,24 @@ import (
 	"github.com/goto/raccoon/services/rest/websocket/connection"
 )
 
-// Pinger is worker that pings the connected peers based on ping interval.
-func Pinger(c chan connection.Conn, size int, PingInterval time.Duration, WriteWaitInterval time.Duration) {
+// Pinger is worker that pings the connected peers based on ping interval. This will also shutdown all connection kept in the cset map
+func Pinger(ctx context.Context, c chan connection.Conn, size int, PingInterval time.Duration, WriteWaitInterval time.Duration) {
 	for i := 0; i < size; i++ {
 		go func() {
 			cSet := make(map[identification.Identifier]connection.Conn)
 			ticker := time.NewTicker(PingInterval)
+			defer ticker.Stop()
 			for {
 				select {
+				//close the connection
+				case <-ctx.Done():
+					// shutdown signal received
+					logger.Infof("[websocket.pinger] shutting down, closing %d active connections", len(cSet))
+					for _, conn := range cSet {
+						conn.Close()
+					}
+					logger.Info("[websocket.pinger] - shutting down, successful")
+					return
 				case conn := <-c:
 					cSet[conn.Identifier] = conn
 				case <-ticker.C:

services/rest/websocket/pinger_test.go

@@ -0,0 +1,98 @@
+package websocket
+
+import (
+	"context"
+	"github.com/gorilla/websocket"
+	"github.com/goto/raccoon/services/rest/websocket/connection"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestPinger(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	connCh := make(chan connection.Conn, 1)
+
+	upgrader := connection.NewUpgrader(connection.UpgraderConfig{
+		ReadBufferSize:    1024,
+		WriteBufferSize:   1024,
+		CheckOrigin:       true,
+		MaxUser:           10,
+		PongWaitInterval:  1 * time.Second,
+		WriteWaitInterval: 1 * time.Second,
+		ConnIDHeader:      "X-Conn-ID",
+		ConnGroupHeader:   "X-Conn-Group",
+		ConnGroupDefault:  "default",
+	})
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r)
+		if err != nil {
+			t.Errorf("failed to upgrade: %v", err)
+			return
+		}
+		connCh <- conn
+	}))
+	defer srv.Close()
+
+	wsURL := "ws" + strings.TrimPrefix(srv.URL, "http")
+	client, _, err := websocket.DefaultDialer.Dial(wsURL, http.Header{
+		"X-Conn-ID":    {"1"},
+		"X-Conn-Group": {"test"},
+	})
+	if err != nil {
+		t.Fatalf("failed to dial test server: %v", err)
+	}
+
+	// ping detection
+	pingReceived := make(chan struct{}, 1)
+	client.SetPingHandler(func(appData string) error {
+		t.Logf("client received ping: %s", appData)
+		select { // avoid blocking if already signaled
+		case pingReceived <- struct{}{}:
+		default:
+		}
+		return client.WriteControl(websocket.PongMessage, []byte(appData), time.Now().Add(time.Second))
+	})
+
+	// close detection
+	closed := make(chan struct{})
+	var once sync.Once
+	go func() {
+		for {
+			_, _, err := client.ReadMessage()
+			if err != nil {
+				once.Do(func() { close(closed) })
+				return
+			}
+		}
+	}()
+
+	// start pinger
+	go Pinger(ctx, connCh, 1, 20*time.Millisecond, 5*time.Millisecond)
+
+	// assert ping received
+	select {
+	case <-pingReceived:
+		t.Log("client received ping")
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("client did not receive ping")
+	}
+
+	// now cancel -> triggers connection close
+	cancel()
+
+	select {
+	case <-closed:
+		t.Log("connection closed after cancel()")
+	case <-time.After(500 * time.Millisecond):
+		t.Fatal("connection not closed after cancel()")
+	}
+
+	// close client at the very end
+	client.Close()
+}

services/services.go

@@ -43,13 +43,13 @@ func (s *Services) Shutdown(ctx context.Context) {
 	}
 }
 
-func Create(b chan collection.CollectRequest) Services {
+func Create(b chan collection.CollectRequest, ctx context.Context) Services {
 	c := collection.NewChannelCollector(b)
 	return Services{
 		b: []bootstrapper{
 			grpc.NewGRPCService(c),
 			pprof.NewPprofService(),
-			rest.NewRestService(c),
+			rest.NewRestService(c, ctx),
 		},
 	}
 }