vnet: add exponential backoff for failed reconnections (#5035)

pkg/util/net/conn.go

@@ -135,11 +135,11 @@ type CloseNotifyConn struct {
 	// 1 means closed
 	closeFlag int32
 
-	closeFn func()
+	closeFn func(error)
 }
 
-// closeFn will be only called once
-func WrapCloseNotifyConn(c net.Conn, closeFn func()) net.Conn {
+// closeFn will be only called once with the error (nil if Close() was called, non-nil if CloseWithError() was called)
+func WrapCloseNotifyConn(c net.Conn, closeFn func(error)) *CloseNotifyConn {
 	return &CloseNotifyConn{
 		Conn:    c,
 		closeFn: closeFn,
@@ -151,12 +151,25 @@ func (cc *CloseNotifyConn) Close() (err error) {
 	if pflag == 0 {
 		err = cc.Conn.Close()
 		if cc.closeFn != nil {
-			cc.closeFn()
+			cc.closeFn(nil)
 		}
 	}
 	return
 }
 
+// CloseWithError closes the connection and passes the error to the close callback.
+func (cc *CloseNotifyConn) CloseWithError(err error) error {
+	pflag := atomic.SwapInt32(&cc.closeFlag, 1)
+	if pflag == 0 {
+		closeErr := cc.Conn.Close()
+		if cc.closeFn != nil {
+			cc.closeFn(err)
+		}
+		return closeErr
+	}
+	return nil
+}
+
 type StatsConn struct {
 	net.Conn
 

pkg/util/net/websocket.go

@@ -32,7 +32,7 @@ func NewWebsocketListener(ln net.Listener) (wl *WebsocketListener) {
 	muxer := http.NewServeMux()
 	muxer.Handle(FrpWebsocketPath, websocket.Handler(func(c *websocket.Conn) {
 		notifyCh := make(chan struct{})
-		conn := WrapCloseNotifyConn(c, func() {
+		conn := WrapCloseNotifyConn(c, func(_ error) {
 			close(notifyCh)
 		})
 		wl.acceptCh <- conn