From f2680ad45fc4118cb9604749992d7d5863b6a7d7 Mon Sep 17 00:00:00 2001 From: Jordan Krage Date: Mon, 10 Jun 2024 12:58:42 -0500 Subject: [PATCH] core/capabilities/remote/target: server.Receive - cancel Context; use WaitGroup (#13482) --- core/capabilities/remote/target/server.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/core/capabilities/remote/target/server.go b/core/capabilities/remote/target/server.go index 8ccbecf530e..83451049e75 100644 --- a/core/capabilities/remote/target/server.go +++ b/core/capabilities/remote/target/server.go @@ -111,7 +111,6 @@ func (r *server) expireRequests() { func (r *server) Receive(msg *types.MessageBody) { r.receiveLock.Lock() defer r.receiveLock.Unlock() - ctx, _ := r.stopCh.NewCtx() if msg.Method != types.MethodExecute { r.lggr.Errorw("received request for unsupported method type", "method", msg.Method) @@ -137,7 +136,11 @@ func (r *server) Receive(msg *types.MessageBody) { req := r.requestIDToRequest[requestID] + r.wg.Add(1) go func() { + defer r.wg.Done() + ctx, cancel := r.stopCh.NewCtx() + defer cancel() err := req.OnMessage(ctx, msg) if err != nil { r.lggr.Errorw("request failed to OnMessage new message", "request", req, "err", err)