From 3e303bc36cec8d8ab142bf4b328941da1362051e Mon Sep 17 00:00:00 2001 From: Vamsi Mundra Date: Fri, 6 Dec 2024 10:42:09 -0800 Subject: [PATCH] add stop function to ContextWServer --- cmdutil/server.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cmdutil/server.go b/cmdutil/server.go index dbf6325b..77d22b3f 100644 --- a/cmdutil/server.go +++ b/cmdutil/server.go @@ -58,14 +58,17 @@ func (sf ServerFuncs) Stop(err error) { // NewContextServer that when Run(), calls the given function with a context // that is canceled when Stop() is called. -func NewContextServer(fn func(context.Context) error) Server { +func NewContextServer(fn func(context.Context) error, stopFunc ...func(error)) Server { ctx, cancel := context.WithCancel(context.Background()) return &ServerFuncs{ RunFunc: func() error { return fn(ctx) }, - StopFunc: func(error) { + StopFunc: func(err error) { + for _, stop := range stopFunc { + stop(err) + } cancel() }, }