From 2fc3320e59d4015027182577587c5f42002fc43e Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Thu, 10 Dec 2015 22:42:17 -0800
Subject: [PATCH] rafthttp: kill the receiving body timeout TODO in snapshot
 sender

---
 rafthttp/snapshot_sender.go | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/rafthttp/snapshot_sender.go b/rafthttp/snapshot_sender.go
index 8a9a99e93..53e348278 100644
--- a/rafthttp/snapshot_sender.go
+++ b/rafthttp/snapshot_sender.go
@@ -27,6 +27,11 @@ import (
 	"github.com/coreos/etcd/snap"
 )
 
+var (
+	// timeout for reading snapshot response body
+	snapResponseReadTimeout = 5 * time.Second
+)
+
 type snapshotSender struct {
 	from, to types.ID
 	cid      types.ID
@@ -108,9 +113,6 @@ func (s *snapshotSender) post(req *http.Request) (err error) {
 	result := make(chan responseAndError, 1)
 
 	go func() {
-		// TODO: cancel the request if it has waited for a long time(~5s) after
-		// it has write out the full request body, which helps to avoid receiver
-		// dies when sender is waiting for response
 		// TODO: the snapshot could be large and eat up all resources when writing
 		// it out. Send it block by block and rest some time between to give the
 		// time for main loop to run.
@@ -119,8 +121,12 @@ func (s *snapshotSender) post(req *http.Request) (err error) {
 			result <- responseAndError{resp, nil, err}
 			return
 		}
+
+		// close the response body when timeouts.
+		// prevents from reading the body forever when the other side dies right after
+		// successfully receives the request body.
+		time.AfterFunc(snapResponseReadTimeout, func() { resp.Body.Close() })
 		body, err := ioutil.ReadAll(resp.Body)
-		resp.Body.Close()
 		result <- responseAndError{resp, body, err}
 	}()