From 3f45cfd6e86cc03af1af1fdeda3336c79a6d9310 Mon Sep 17 00:00:00 2001
From: neza2017 <yefu.chen@zilliz.com>
Date: Mon, 1 Feb 2021 11:44:02 +0800
Subject: [PATCH] Fix data node

Signed-off-by: neza2017 <yefu.chen@zilliz.com>
---
 internal/datanode/data_node.go | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/internal/datanode/data_node.go b/internal/datanode/data_node.go
index 9d834b7910..ca55446b5d 100644
--- a/internal/datanode/data_node.go
+++ b/internal/datanode/data_node.go
@@ -56,6 +56,7 @@ type (
 
 	DataNode struct {
 		ctx    context.Context
+		cancel context.CancelFunc
 		NodeID UniqueID
 		Role   string
 		State  internalpb2.StateCode
@@ -77,8 +78,10 @@ type (
 func NewDataNode(ctx context.Context) *DataNode {
 
 	Params.Init()
+	ctx2, cancel2 := context.WithCancel(ctx)
 	node := &DataNode{
-		ctx:             ctx,
+		ctx:             ctx2,
+		cancel:          cancel2,
 		NodeID:          Params.NodeID, // GOOSE TODO How to init
 		Role:            typeutil.DataNodeRole,
 		State:           internalpb2.StateCode_INITIALIZING,
@@ -216,7 +219,7 @@ func (node *DataNode) FlushSegments(in *datapb.FlushSegRequest) error {
 }
 
 func (node *DataNode) Stop() error {
-	<-node.ctx.Done()
+	node.cancel()
 
 	// close services
 	if node.dataSyncService != nil {