diff --git a/code/client/pool/conn.go b/code/client/pool/conn.go index 17d786a..a70f77f 100644 --- a/code/client/pool/conn.go +++ b/code/client/pool/conn.go @@ -2,6 +2,7 @@ package pool import ( "context" + "errors" "natpass/code/network" "strings" "sync" @@ -126,6 +127,20 @@ func (conn *Conn) loopRead(cancel context.CancelFunc) { } } +func loopWrite(conn *network.Conn, msg *network.Msg, timeout time.Duration) error { + var err error + for i := 0; i < 10; i++ { + err = conn.WriteMessage(msg, timeout) + if err != nil { + if strings.Contains(err.Error(), "i/o timeout") { + continue + } + return err + } + } + return errors.New("too many retry") +} + func (conn *Conn) loopWrite(cancel context.CancelFunc) { defer conn.Close() defer cancel() @@ -135,12 +150,11 @@ func (conn *Conn) loopWrite(cancel context.CancelFunc) { return } msg.From = conn.ID - err := conn.conn.WriteMessage(msg, conn.parent.cfg.WriteTimeout) - if err == nil { - continue + err := loopWrite(conn.conn, msg, conn.parent.cfg.WriteTimeout) + if err != nil { + logging.Error("write message error on %s: %v", conn.ID, err) + return } - logging.Error("write message error on %s: %v", conn.ID, err) - return } } diff --git a/code/server/handler/client.go b/code/server/handler/client.go index 0844d1a..4c4f7ff 100644 --- a/code/server/handler/client.go +++ b/code/server/handler/client.go @@ -1,6 +1,7 @@ package handler import ( + "errors" "natpass/code/network" "strings" "sync" @@ -57,7 +58,16 @@ func (c *client) run() { } func (c *client) writeMessage(msg *network.Msg) error { - return c.c.WriteMessage(msg, c.parent.cfg.WriteTimeout) + for i := 0; i < 10; i++ { + err := c.c.WriteMessage(msg, c.parent.cfg.WriteTimeout) + if err != nil { + if strings.Contains(err.Error(), "i/o timeout") { + continue + } + return err + } + } + return errors.New("too many retry") } func (c *client) addLink(id string) {