From eadbac42afdac88b9afdbd181e95e0d18209736d Mon Sep 17 00:00:00 2001
From: gagliardetto <gagliardetto@users.noreply.github.com>
Date: Sat, 19 Aug 2023 17:04:08 +0200
Subject: [PATCH] Improve closing

---
 cmd-rpc.go    |  4 ++--
 multiepoch.go | 18 ++++++++++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/cmd-rpc.go b/cmd-rpc.go
index 7df17b65..fe7abdb9 100644
--- a/cmd-rpc.go
+++ b/cmd-rpc.go
@@ -142,7 +142,7 @@ func newCmd_rpc() *cli.Command {
 				klog.Infof("Found %d directories; will start watching them for changes ...", len(dirs))
 				spew.Dump(dirs)
 
-				ctx, cancel := context.WithCancel(context.Background())
+				ctx, cancel := context.WithCancel(c.Context)
 				defer cancel()
 
 				err = onFileChanged(ctx, dirs, func(event fsnotify.Event) {
@@ -234,7 +234,7 @@ func newCmd_rpc() *cli.Command {
 				}
 			}
 
-			return multi.ListenAndServe(listenOn, listenerConfig)
+			return multi.ListenAndServe(c.Context, listenOn, listenerConfig)
 		},
 	}
 }
diff --git a/multiepoch.go b/multiepoch.go
index 2a396b47..e31a6497 100644
--- a/multiepoch.go
+++ b/multiepoch.go
@@ -175,7 +175,7 @@ func LoadProxyConfig(configFilepath string) (*ProxyConfig, error) {
 }
 
 // ListeAndServe starts listening on the configured address and serves the RPC API.
-func (m *MultiEpoch) ListenAndServe(listenOn string, lsConf *ListenerConfig) error {
+func (m *MultiEpoch) ListenAndServe(ctx context.Context, listenOn string, lsConf *ListenerConfig) error {
 	handler := newMultiEpochHandler(m, lsConf)
 	handler = fasthttp.CompressHandler(handler)
 
@@ -206,7 +206,21 @@ func (m *MultiEpoch) ListenAndServe(listenOn string, lsConf *ListenerConfig) err
 	}
 
 	klog.Infof("RPC server listening on %s", listenOn)
-	return fasthttp.ListenAndServe(listenOn, handler)
+
+	s := &fasthttp.Server{
+		Handler:            handler,
+		MaxRequestBodySize: 1024 * 1024,
+	}
+	go func() {
+		// listen for context cancellation
+		<-ctx.Done()
+		klog.Info("RPC server shutting down...")
+		defer klog.Info("RPC server shut down")
+		if err := s.ShutdownWithContext(ctx); err != nil {
+			klog.Errorf("Error while shutting down RPC server: %s", err)
+		}
+	}()
+	return s.ListenAndServe(listenOn)
 }
 
 func randomRequestID() string {