From a33d6ae4fb22a4d72d714733b3045272e050c0d7 Mon Sep 17 00:00:00 2001
From: xingweizheng 00591739 <xingweizheng@huawei.com>
Date: Mon, 9 Nov 2020 13:54:57 +0800
Subject: [PATCH] fix panic when user knock ctrl+c when pull, push and save

---
 daemon/pull.go | 24 +++++++++---------------
 daemon/push.go | 24 +++++++++---------------
 daemon/save.go | 29 +++++++++--------------------
 3 files changed, 27 insertions(+), 50 deletions(-)

diff --git a/daemon/pull.go b/daemon/pull.go
index f9dee3e..56be755 100644
--- a/daemon/pull.go
+++ b/daemon/pull.go
@@ -59,23 +59,17 @@ func (b *Backend) Pull(req *pb.PullRequest, stream pb.Control_PullServer) error
 	eg.Go(pullMessageHandler(stream, opt.logger))
 	errC := make(chan error, 1)
 
-	go func() { errC <- eg.Wait() }()
+	errC <- eg.Wait()
 	defer close(errC)
 
-	select {
-	case err2 := <-errC:
-		if err2 != nil {
-			return err2
-		}
-	case _, ok := <-stream.Context().Done():
-		if !ok {
-			logrus.WithField(util.LogKeySessionID, opt.pullID).Info("Channel stream done closed")
-			return nil
-		}
-		err := egCtx.Err()
-		if err != nil && err != context.Canceled {
-			logrus.WithField(util.LogKeySessionID, opt.pullID).Warnf("Stream closed with: %v", err)
-		}
+	err, ok := <-errC
+	if !ok {
+		logrus.WithField(util.LogKeySessionID, opt.pullID).Info("Channel errC closed")
+		return nil
+	}
+	if err != nil {
+		logrus.WithField(util.LogKeySessionID, opt.pullID).Warnf("Stream closed with: %v", err)
+		return err
 	}
 
 	return nil
diff --git a/daemon/push.go b/daemon/push.go
index 712062e..ea5e47c 100644
--- a/daemon/push.go
+++ b/daemon/push.go
@@ -63,23 +63,17 @@ func (b *Backend) Push(req *pb.PushRequest, stream pb.Control_PushServer) error
 	eg.Go(pushMessageHandler(stream, opt.logger))
 	errC := make(chan error, 1)
 
-	go func() { errC <- eg.Wait() }()
+	errC <- eg.Wait()
 	defer close(errC)
 
-	select {
-	case err2 := <-errC:
-		if err2 != nil {
-			return err2
-		}
-	case _, ok := <-stream.Context().Done():
-		if !ok {
-			logrus.WithField(util.LogKeySessionID, opt.pushID).Info("Channel stream done closed")
-			return nil
-		}
-		err := egCtx.Err()
-		if err != nil && err != context.Canceled {
-			logrus.WithField(util.LogKeySessionID, opt.pushID).Warnf("Stream closed with: %v", err)
-		}
+	err, ok := <-errC
+	if !ok {
+		logrus.WithField(util.LogKeySessionID, opt.pushID).Info("Channel errC closed")
+		return nil
+	}
+	if err != nil {
+		logrus.WithField(util.LogKeySessionID, opt.pushID).Warnf("Stream closed with: %v", err)
+		return err
 	}
 
 	return nil
diff --git a/daemon/save.go b/daemon/save.go
index 13ca8cd..156a1c2 100644
--- a/daemon/save.go
+++ b/daemon/save.go
@@ -102,33 +102,22 @@ func (b *Backend) Save(req *pb.SaveRequest, stream pb.Control_SaveServer) (err e
 	}
 
 	ctx := context.WithValue(stream.Context(), util.LogFieldKey(util.LogKeySessionID), opts.saveID)
-	eg, egCtx := errgroup.WithContext(ctx)
+	eg, _ := errgroup.WithContext(ctx)
 
 	eg.Go(exportHandler(ctx, stream, opts))
 	eg.Go(messageHandler(stream, opts.logger))
 	errC := make(chan error, 1)
 
-	go func() { errC <- eg.Wait() }()
+	errC <- eg.Wait()
 	defer close(errC)
 
-	select {
-	case err, ok = <-errC:
-		if !ok {
-			opts.logEntry.Info("Channel errC closed")
-			return nil
-		}
-		if err != nil {
-			return err
-		}
-	case _, ok := <-stream.Context().Done():
-		if !ok {
-			opts.logEntry.Info("Channel stream done closed")
-			return nil
-		}
-		err = egCtx.Err()
-		if err != nil && err != context.Canceled {
-			opts.logEntry.Infof("Stream closed with: %v", err)
-		}
+	err, ok = <-errC
+	if !ok {
+		opts.logEntry.Info("Channel errC closed")
+		return nil
+	}
+	if err != nil {
+		return err
 	}
 
 	return nil
-- 
1.8.3.1