From 02f4ae6c56474b1f4e747916812b38134d503349 Mon Sep 17 00:00:00 2001
From: Derek <crq@kernel.org>
Date: Thu, 22 May 2014 23:58:56 -0700
Subject: [PATCH] Use Timeout Conn wrapper to set read deadline for downloading
 layer

Docker-DCO-1.1-Signed-off-by: Derek <crq@kernel.org> (github: crquan)
---
 registry/registry.go | 11 +++++++++++
 server/server.go     | 44 +++++++++++++++++++++++++++++++++-----------
 utils/timeoutconn.go | 26 ++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 11 deletions(-)
 create mode 100644 utils/timeoutconn.go

diff --git a/registry/registry.go b/registry/registry.go
index 2e3e7e03a7..3d0a3ed2da 100644
--- a/registry/registry.go
+++ b/registry/registry.go
@@ -726,7 +726,17 @@ type Registry struct {
 }
 
 func NewRegistry(authConfig *AuthConfig, factory *utils.HTTPRequestFactory, indexEndpoint string) (r *Registry, err error) {
+	httpDial := func(proto string, addr string) (net.Conn, error) {
+		conn, err := net.Dial(proto, addr)
+		if err != nil {
+			return nil, err
+		}
+		conn = utils.NewTimeoutConn(conn, time.Duration(1)*time.Minute)
+		return conn, nil
+	}
+
 	httpTransport := &http.Transport{
+		Dial:              httpDial,
 		DisableKeepAlives: true,
 		Proxy:             http.ProxyFromEnvironment,
 	}
@@ -738,6 +748,7 @@ func NewRegistry(authConfig *AuthConfig, factory *utils.HTTPRequestFactory, inde
 		},
 		indexEndpoint: indexEndpoint,
 	}
+
 	r.client.Jar, err = cookiejar.New(nil)
 	if err != nil {
 		return nil, err
diff --git a/server/server.go b/server/server.go
index 6d398807bd..0deda0d955 100644
--- a/server/server.go
+++ b/server/server.go
@@ -27,6 +27,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net"
 	"net/http"
 	"net/url"
 	"os"
@@ -1134,17 +1135,38 @@ func (srv *Server) pullImage(r *registry.Registry, out io.Writer, imgID, endpoin
 				}
 			}
 
-			// Get the layer
-			out.Write(sf.FormatProgress(utils.TruncateID(id), "Pulling fs layer", nil))
-			layer, err := r.GetRemoteImageLayer(img.ID, endpoint, token)
-			if err != nil {
-				out.Write(sf.FormatProgress(utils.TruncateID(id), "Error pulling dependent layers", nil))
-				return err
-			}
-			defer layer.Close()
-			if err := srv.daemon.Graph().Register(imgJSON, utils.ProgressReader(layer, imgSize, out, sf, false, utils.TruncateID(id), "Downloading"), img); err != nil {
-				out.Write(sf.FormatProgress(utils.TruncateID(id), "Error downloading dependent layers", nil))
-				return err
+			for j := 1; j <= retries; j++ {
+				// Get the layer
+				status := "Pulling fs layer"
+				if j > 1 {
+					status = fmt.Sprintf("Pulling fs layer [retries: %d]", j)
+				}
+				out.Write(sf.FormatProgress(utils.TruncateID(id), status, nil))
+				layer, err := r.GetRemoteImageLayer(img.ID, endpoint, token)
+				if uerr, ok := err.(*url.Error); ok {
+					err = uerr.Err
+				}
+				if terr, ok := err.(net.Error); ok && terr.Timeout() && j < retries {
+					time.Sleep(time.Duration(j) * 500 * time.Millisecond)
+					continue
+				} else if err != nil {
+					out.Write(sf.FormatProgress(utils.TruncateID(id), "Error pulling dependent layers", nil))
+					return err
+				}
+				defer layer.Close()
+
+				err = srv.daemon.Graph().Register(imgJSON,
+					utils.ProgressReader(layer, imgSize, out, sf, false, utils.TruncateID(id), "Downloading"),
+					img)
+				if terr, ok := err.(net.Error); ok && terr.Timeout() && j < retries {
+					time.Sleep(time.Duration(j) * 500 * time.Millisecond)
+					continue
+				} else if err != nil {
+					out.Write(sf.FormatProgress(utils.TruncateID(id), "Error downloading dependent layers", nil))
+					return err
+				} else {
+					break
+				}
 			}
 		}
 		out.Write(sf.FormatProgress(utils.TruncateID(id), "Download complete", nil))
diff --git a/utils/timeoutconn.go b/utils/timeoutconn.go
new file mode 100644
index 0000000000..a3231c7ee3
--- /dev/null
+++ b/utils/timeoutconn.go
@@ -0,0 +1,26 @@
+package utils
+
+import (
+	"net"
+	"time"
+)
+
+func NewTimeoutConn(conn net.Conn, timeout time.Duration) net.Conn {
+	return &TimeoutConn{conn, timeout}
+}
+
+// A net.Conn that sets a deadline for every Read or Write operation
+type TimeoutConn struct {
+	net.Conn
+	timeout time.Duration
+}
+
+func (c *TimeoutConn) Read(b []byte) (int, error) {
+	if c.timeout > 0 {
+		err := c.Conn.SetReadDeadline(time.Now().Add(c.timeout))
+		if err != nil {
+			return 0, err
+		}
+	}
+	return c.Conn.Read(b)
+}