Force close all open connections at shutdown

Author: altafan

server/internal/interface/grpc/handlers/arkservice.go

@@ -33,7 +33,10 @@ type handler struct {
 	transactionsListenerHandler *listenerHanlder[*arkv1.GetTransactionsStreamResponse]
 	addressSubsHandler          *listenerHanlder[*arkv1.SubscribeForAddressResponse]
 
-	stopCh <-chan struct{}
+	stopCh                  <-chan struct{}
+	stopRoundEventsCh       chan struct{}
+	stopTransactionEventsCh chan struct{}
+	stopAddressEventsCh     chan struct{}
 }
 
 func NewHandler(version string, service application.Service, stopCh <-chan struct{}) service {
@@ -44,8 +47,12 @@ func NewHandler(version string, service application.Service, stopCh <-chan struc
 		transactionsListenerHandler: newListenerHandler[*arkv1.GetTransactionsStreamResponse](),
 		addressSubsHandler:          newListenerHandler[*arkv1.SubscribeForAddressResponse](),
 		stopCh:                      stopCh,
+		stopRoundEventsCh:           make(chan struct{}, 1),
+		stopTransactionEventsCh:     make(chan struct{}, 1),
+		stopAddressEventsCh:         make(chan struct{}, 1),
 	}
 
+	go h.listenToStop()
 	go h.listenToEvents()
 	go h.listenToTxEvents()
 
@@ -304,7 +311,7 @@ func (h *handler) GetEventStream(
 
 	for {
 		select {
-		case <-h.stopCh:
+		case <-h.stopRoundEventsCh:
 			return nil
 		case <-stream.Context().Done():
 			return nil
@@ -456,6 +463,8 @@ func (h *handler) GetTransactionsStream(
 
 	for {
 		select {
+		case <-h.stopTransactionEventsCh:
+			return nil
 		case <-stream.Context().Done():
 			return nil
 		case ev := <-listener.ch:
@@ -524,6 +533,8 @@ func (h *handler) SubscribeForAddress(
 
 	for {
 		select {
+		case <-h.stopAddressEventsCh:
+			return nil
 		case <-stream.Context().Done():
 			return nil
 		case ev := <-listener.ch:
@@ -534,6 +545,14 @@ func (h *handler) SubscribeForAddress(
 	}
 }
 
+func (h *handler) listenToStop() {
+	<-h.stopCh
+	h.stopRoundEventsCh <- struct{}{}
+	h.stopTransactionEventsCh <- struct{}{}
+	h.stopAddressEventsCh <- struct{}{}
+
+}
+
 // listenToEvents forwards events from the application layer to the set of listeners
 func (h *handler) listenToEvents() {
 	channel := h.svc.GetEventsChannel(context.Background())

server/internal/interface/grpc/service.go

@@ -4,6 +4,10 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"net/http"
+	"path/filepath"
+	"strings"
+
 	arkv1 "github.com/ark-network/ark/api-spec/protobuf/gen/ark/v1"
 	"github.com/ark-network/ark/server/internal/config"
 	"github.com/ark-network/ark/server/internal/core/application"
@@ -23,9 +27,6 @@ import (
 	"google.golang.org/grpc/credentials/insecure"
 	grpchealth "google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/protobuf/encoding/protojson"
-	"net/http"
-	"path/filepath"
-	"strings"
 )
 
 const (
@@ -151,9 +152,6 @@ func (s *service) start(withAppSvc bool) error {
 }
 
 func (s *service) stop(withAppSvc bool) {
-	//nolint:all
-	s.server.Shutdown(context.Background())
-	log.Info("stopped grpc server")
 	if withAppSvc {
 		s.stopCh <- struct{}{}
 		appSvc, _ := s.appConfig.AppService()
@@ -162,6 +160,9 @@ func (s *service) stop(withAppSvc bool) {
 			log.Info("stopped app service")
 		}
 	}
+	//nolint:all
+	s.server.Shutdown(context.Background())
+	log.Info("stopped grpc server")
 }
 
 func (s *service) newServer(tlsConfig *tls.Config, withAppSvc bool) error {