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 {