From fbc98fc65cec22f7237d4a86b704e47b4b3ccf81 Mon Sep 17 00:00:00 2001 From: Anil Vishnoi Date: Sat, 27 Apr 2024 01:46:35 -0700 Subject: [PATCH] Add NotifyContext and pass it to http request to abort the request of the worker is killed, so worker won't be stuck on http request Signed-off-by: Anil Vishnoi --- worker/cmd/generate.go | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/worker/cmd/generate.go b/worker/cmd/generate.go index 62820ab2..b6f132d8 100644 --- a/worker/cmd/generate.go +++ b/worker/cmd/generate.go @@ -132,7 +132,9 @@ var generateCmd = &cobra.Command{ Run: func(cmd *cobra.Command, args []string) { logger := initLogger(Debug) sugar := logger.Sugar() - ctx := cmd.Context() + + ctx, cancel := signal.NotifyContext(cmd.Context(), syscall.SIGTERM, syscall.SIGQUIT, syscall.SIGINT) + defer cancel() sugar.Info("Starting generate worker") @@ -916,7 +918,7 @@ func (w *Worker) datagenSvc(taxonomyFiles []string, outputDir string, numSamples requestURL = strings.Replace(requestURL, "skill", "knowledge", -1) } - request, err := http.NewRequest("POST", requestURL, bytes.NewBuffer(jsonData)) + request, err := http.NewRequestWithContext(w.ctx, "POST", requestURL, bytes.NewBuffer(jsonData)) if err != nil { return nil, fmt.Errorf("failed to create request: %w", err) }