websocket: optimize async close delay

lesismal · Sep 14, 2023 · 5ca5456 · 5ca5456
1 parent eb81029
commit 5ca5456
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
@@ -100,7 +100,7 @@ func (c *Conn) Close() error {
 		return nil
 	}
 	if c.IsAsyncWrite() {
-		c.Engine.AfterFunc(time.Second, func() { c.Conn.Close() })
+		c.Engine.AfterFunc(c.BlockingModAsyncCloseDelay, func() { c.Conn.Close() })
 		return nil
 	}
 	return c.Conn.Close()
@@ -780,6 +780,9 @@ func NewConn(u *Upgrader, c net.Conn, subprotocol string, remoteCompressionEnabl
 	if asyncWrite {
 		wsc.sendQueue = make([][]byte, u.BlockingModSendQueueInitSize)[:0]
 		wsc.sendQueueSize = u.BlockingModSendQueueMaxSize
+		if wsc.BlockingModAsyncCloseDelay <= 0 {
+			wsc.BlockingModAsyncCloseDelay = DefaultBlockingModAsyncCloseDelay
+		}
 	}
 	return wsc
 }

diff --git a/nbhttp/websocket/upgrader.go b/nbhttp/websocket/upgrader.go
@@ -34,16 +34,19 @@ var (
 
 	DefaultBlockingModSendQueueMaxSize = 0
 
+	DefaultBlockingModAsyncCloseDelay = time.Second / 10
+
 	// DefaultEngine will be set to a Upgrader.Engine to handle details such as buffers.
 	DefaultEngine = nbhttp.NewEngine(nbhttp.Config{
 		ReleaseWebsocketPayload: true,
 	})
 )
 
 type commonFields struct {
-	Engine             *nbhttp.Engine
-	KeepaliveTime      time.Duration
-	MessageLengthLimit int
+	Engine                     *nbhttp.Engine
+	KeepaliveTime              time.Duration
+	MessageLengthLimit         int
+	BlockingModAsyncCloseDelay time.Duration
 
 	enableCompression bool
 	compressionLevel  int
@@ -92,8 +95,9 @@ type Upgrader struct {
 func NewUpgrader() *Upgrader {
 	u := &Upgrader{
 		commonFields: commonFields{
-			Engine:           DefaultEngine,
-			compressionLevel: defaultCompressionLevel,
+			Engine:                     DefaultEngine,
+			compressionLevel:           defaultCompressionLevel,
+			BlockingModAsyncCloseDelay: DefaultBlockingModAsyncCloseDelay,
 		},
 		BlockingModReadBufferSize:    DefaultBlockingReadBufferSize,
 		BlockingModAsyncWrite:        DefaultBlockingModAsyncWrite,