[INLONG-8631][SDK] Handle context.Done() in Send() (#8632)

Co-authored-by: gunli <[email protected]>
apache · Aug 6, 2023 · eaa1a04 · eaa1a04
1 parent 8012151
commit eaa1a04
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/inlong-sdk/dataproxy-sdk-twins/dataproxy-sdk-golang/dataproxy/worker.go b/inlong-sdk/dataproxy-sdk-twins/dataproxy-sdk-golang/dataproxy/worker.go
@@ -301,8 +301,12 @@ func (w *worker) send(ctx context.Context, msg Message) error {
 	}, true)
 
 	// wait for send done
-	<-doneCh
-	return err
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-doneCh:
+		return err
+	}
 }
 
 func (w *worker) sendAsync(ctx context.Context, msg Message, callback Callback) {