From a709ac2dc5793fdb91f1dea23e5f2419526eb2a6 Mon Sep 17 00:00:00 2001
From: Bilal <bilal@saburly.com>
Date: Wed, 13 May 2020 18:23:56 +0200
Subject: [PATCH] Reconnect after waiting too long for request

---
 config/config.go             |  2 ++
 structures/structures.go     |  1 +
 workerclient/workerclient.go | 30 +++++++++++++++++++++++++++++-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/config/config.go b/config/config.go
index b900027..38a5621 100644
--- a/config/config.go
+++ b/config/config.go
@@ -39,6 +39,7 @@ func loadEnvVariables() {
 func generateClientConfigObject() {
 	ClientConfig.ConnectionsCount = getInt("CLIENT_CONNECTIONS_COUNT")
 	ClientConfig.ConnectionTimeout = getInt("CLIENT_CONNECTION_TIMEOUT")
+	ClientConfig.WaitingTimeout = getInt("CLIENT_WAITING_TIMEOUT")
 	ClientConfig.WorkerServerAddress = getString("WORKER_SERVER_ADDRESS")
 	ClientConfig.RequestMessagePrefix = getString("REQUEST_MESSAGE_PREFIX")
 	ClientConfig.ProxyListBaseURL = getString("PROXY_LIST_BASE_URL")
@@ -57,6 +58,7 @@ func generateServerConfigObject() {
 func initClientConfigDefaultValues() {
 	defaultClientConfigValues["CLIENT_CONNECTIONS_COUNT"] = "5"
 	defaultClientConfigValues["CLIENT_CONNECTION_TIMEOUT"] = "2"
+	defaultClientConfigValues["CLIENT_WAITING_TIMEOUT"] = "60"
 	defaultClientConfigValues["WORKER_SERVER_ADDRESS"] = "127.0.0.1:1338"
 	defaultClientConfigValues["REQUEST_MESSAGE_PREFIX"] = "URL "
 	defaultClientConfigValues["PROXY_LIST_BASE_URL"] = "https://www.proxy-list.download/api/v1/get?type="
diff --git a/structures/structures.go b/structures/structures.go
index 17605a0..e9b809a 100644
--- a/structures/structures.go
+++ b/structures/structures.go
@@ -36,6 +36,7 @@ type WorkerServerConfig struct {
 type ClientConfig struct {
 	ConnectionsCount     int
 	ConnectionTimeout    int // In seconds
+	WaitingTimeout       int // In seconds
 	WorkerServerAddress  string
 	RequestMessagePrefix string
 	ProxyListBaseURL     string
diff --git a/workerclient/workerclient.go b/workerclient/workerclient.go
index b82f4f7..5d548ab 100644
--- a/workerclient/workerclient.go
+++ b/workerclient/workerclient.go
@@ -45,7 +45,35 @@ func startSingleConnection(connectionId int) {
 		}
 
 		for {
-			encodedRequestMessage, err := bufio.NewReader(conn).ReadString('\n')
+			timeout := make(chan bool, 1)
+			go func() {
+				time.Sleep(time.Duration(c.ClientConfig.WaitingTimeout) * time.Second)
+				timeout <- true
+			}()
+
+			requestChann := make(chan bool, 1)
+			var err error
+			var encodedRequestMessage string
+			go func() {
+				encodedRequestMessage, err = bufio.NewReader(conn).ReadString('\n')
+				requestChann <- true
+			}()
+
+			timeoutConnection := false
+
+			select {
+			case <-requestChann:
+				timeoutConnection = false
+			case <-timeout:
+				timeoutConnection = true
+			}
+
+			if timeoutConnection {
+				log.Printf("(%d) Server not sending requests for too long, closing connection\n", connectionId)
+				_ = conn.Close()
+				break
+			}
+
 			if err != nil {
 				log.Printf("(%d) Error receiving request from load balancer server : %s\n", connectionId, err)
 				_ = conn.Close()