From b7cdfaed366b43949821bf71e6f333127509f27f Mon Sep 17 00:00:00 2001
From: fupan <lifupan@gmail.com>
Date: Mon, 15 Oct 2018 16:30:16 +0800
Subject: [PATCH] agent: Fix the issue of stdout hang on builtin proxy

Since the Yamux's keepalive has been disabled both
on the server and the client side, and this brings
a weird issue where the communication between the
proxy and the agent hangs. The same issue has been
fixed in kata proxy by:
"https://github.com/kata-containers/proxy/pull/91".

This commit just cherry-pick that patch here to
fix the same issue on kata builtin proxy.

Fixes: #396

Signed-off-by: fupan <lifupan@gmail.com>
---
 protocols/client/client.go | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/protocols/client/client.go b/protocols/client/client.go
index 451beab07f..a0431301dc 100644
--- a/protocols/client/client.go
+++ b/protocols/client/client.go
@@ -165,6 +165,26 @@ func parse(sock string) (string, *url.URL, error) {
 	return grpcAddr, addr, nil
 }
 
+// This function is meant to run in a go routine since it will send ping
+// commands every second. It behaves as a heartbeat to maintain a proper
+// communication state with the Yamux server in the agent.
+func heartBeat(session *yamux.Session) {
+	if session == nil {
+		return
+	}
+
+	for {
+		if session.IsClosed() {
+			break
+		}
+
+		session.Ping()
+
+		// 1 Hz heartbeat
+		time.Sleep(time.Second)
+	}
+}
+
 func agentDialer(addr *url.URL, enableYamux bool) dialer {
 	var d dialer
 	switch addr.Scheme {
@@ -196,11 +216,15 @@ func agentDialer(addr *url.URL, enableYamux bool) dialer {
 		sessionConfig := yamux.DefaultConfig()
 		// Disable keepAlive since we don't know how much time a container can be paused
 		sessionConfig.EnableKeepAlive = false
+		sessionConfig.ConnectionWriteTimeout = time.Second
 		session, err = yamux.Client(conn, sessionConfig)
 		if err != nil {
 			return nil, err
 		}
 
+		// Start the heartbeat in a separate go routine
+		go heartBeat(session)
+
 		var stream net.Conn
 		stream, err = session.Open()
 		if err != nil {