Merge branch 'norace'

lesismal · Oct 11, 2024 · 3a34aa0 · 3a34aa0
2 parents f4cca21 + a40840c
commit 3a34aa0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 8 deletions.
diff --git a/nbhttp/websocket/conn.go b/nbhttp/websocket/conn.go
@@ -77,6 +77,7 @@ type Conn struct {
 	isInReadingLoop          bool
 	expectingFragments       bool
 	compress                 bool
+	releasePayload           bool
 	msgType                  MessageType
 	message                  []byte
 	bytesCached              []byte
@@ -193,13 +194,13 @@ func (c *Conn) handleMessage(opcode MessageType, body []byte) {
 func (c *Conn) handleProtocolMessage(opcode MessageType, body []byte) {
 	if c.isBlockingMod {
 		c.handleWsMessage(opcode, body)
-		if len(body) > 0 && c.ReleasePayload {
+		if len(body) > 0 && c.releasePayload {
 			c.Engine.BodyAllocator.Free(body)
 		}
 	} else {
 		if !c.Execute(func() {
 			c.handleWsMessage(opcode, body)
-			if len(body) > 0 && c.ReleasePayload {
+			if len(body) > 0 && c.releasePayload {
 				c.Engine.BodyAllocator.Free(body)
 			}
 		}) {
@@ -529,7 +530,7 @@ func (c *Conn) Parse(data []byte) error {
 func (c *Conn) OnMessage(h func(*Conn, MessageType, []byte)) {
 	if h != nil {
 		c.messageHandler = func(c *Conn, messageType MessageType, data []byte) {
-			if c.ReleasePayload && len(data) > 0 {
+			if c.releasePayload && len(data) > 0 {
 				defer c.Engine.BodyAllocator.Free(data)
 			}
 			if !c.closed {
@@ -545,7 +546,7 @@ func (c *Conn) OnMessage(h func(*Conn, MessageType, []byte)) {
 func (c *Conn) OnDataFrame(h func(*Conn, MessageType, bool, []byte)) {
 	if h != nil {
 		c.dataFrameHandler = func(c *Conn, messageType MessageType, fin bool, data []byte) {
-			if c.ReleasePayload {
+			if c.releasePayload {
 				defer c.Engine.BodyAllocator.Free(data)
 			}
 			h(c, messageType, fin, data)

diff --git a/nbhttp/websocket/upgrader.go b/nbhttp/websocket/upgrader.go
@@ -243,7 +243,7 @@ func (u *Upgrader) OnOpen(h func(*Conn)) {
 func (u *Upgrader) OnMessage(h func(*Conn, MessageType, []byte)) {
 	if h != nil {
 		u.messageHandler = func(c *Conn, messageType MessageType, data []byte) {
-			if c.ReleasePayload && len(data) > 0 {
+			if c.releasePayload && len(data) > 0 {
 				defer c.Engine.BodyAllocator.Free(data)
 			}
 			if !c.closed {
@@ -259,7 +259,7 @@ func (u *Upgrader) OnMessage(h func(*Conn, MessageType, []byte)) {
 func (u *Upgrader) OnDataFrame(h func(*Conn, MessageType, bool, []byte)) {
 	if h != nil {
 		u.dataFrameHandler = func(c *Conn, messageType MessageType, fin bool, data []byte) {
-			if c.ReleasePayload {
+			if c.releasePayload {
 				defer c.Engine.BodyAllocator.Free(data)
 			}
 			h(c, messageType, fin, data)
@@ -544,8 +544,13 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
 		}
 	}
 
-	if !wsc.ReleasePayload {
-		wsc.ReleasePayload = wsc.Engine.ReleaseWebsocketPayload
+	// If upgrader.ReleasePayload is false, which maybe by default,
+	// then set it to engine.ReleaseWebsocketPayload.
+	// Considering compatibility with old versions, should set it to
+	// upgrader.ReleasePayload only when upgrader.ReleasePayload is true.
+	wsc.releasePayload = u.ReleasePayload
+	if !wsc.releasePayload {
+		wsc.releasePayload = wsc.Engine.ReleaseWebsocketPayload
 	}
 
 	return wsc, nil