From 8cfde43f5d9c4a87eb69597701cccd5fe776116c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Mon, 1 Aug 2022 19:28:27 +0800
Subject: [PATCH] =?UTF-8?q?TCP=E8=BF=9E=E6=8E=A5=E6=BA=90=E7=AB=99?=
 =?UTF-8?q?=E5=A4=B1=E8=B4=A5=E6=97=B6=E5=85=B3=E9=97=AD=E8=BF=9E=E6=8E=A5?=
 =?UTF-8?q?/TCP=E8=BF=9E=E6=8E=A5=E4=B9=9F=E6=94=AF=E6=8C=81=E5=A4=87?=
 =?UTF-8?q?=E7=94=A8=E6=BA=90=E7=AB=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/listener_tcp.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/internal/nodes/listener_tcp.go b/internal/nodes/listener_tcp.go
index 47f7a70..38114ae 100644
--- a/internal/nodes/listener_tcp.go
+++ b/internal/nodes/listener_tcp.go
@@ -124,6 +124,7 @@ func (this *TCPListener) handleConn(conn net.Conn) error {
 
 	originConn, err := this.connectOrigin(server.Id, server.ReverseProxy, conn.RemoteAddr().String())
 	if err != nil {
+		_ = conn.Close()
 		return err
 	}
 
@@ -243,8 +244,19 @@ func (this *TCPListener) connectOrigin(serverId int64, reverseProxy *serverconfi
 		conn, addr, err = OriginConnect(origin, this.port, remoteAddr, requestHost)
 		if err != nil {
 			remotelogs.ServerError(serverId, "TCP_LISTENER", "unable to connect origin server: "+addr+": "+err.Error(), "", nil)
+
+			SharedOriginStateManager.Fail(origin, requestHost, reverseProxy, func() {
+				reverseProxy.ResetScheduling()
+			})
+
 			continue
 		} else {
+			if !origin.IsOk {
+				SharedOriginStateManager.Success(origin, func() {
+					reverseProxy.ResetScheduling()
+				})
+			}
+
 			return
 		}
 	}