gogrlx/nats-server
1
0
Fork 0
mirror of https://github.com/gogrlx/nats-server.git synced 2026-04-02 03:38:42 -07:00
Code Issues Packages Projects Releases Wiki Activity

Reduce allocations in WebSocket compression

Browse Source 
Signed-off-by: Neil Twigg <neil@nats.io>
This commit is contained in:
Neil Twigg
2023-10-04 12:15:05 +01:00
parent dbe700d192
commit 6b65452bc7
2 changed files with 32 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
server/client.go
View File

@@ -346,20 +346,23 @@ func nbPoolGet(sz int) []byte {
}
}
func nbPoolPut(b []byte) {
switch cap(b) {
case nbPoolSizeSmall:
b := (*[nbPoolSizeSmall]byte)(b[0:nbPoolSizeSmall])
nbPoolSmall.Put(b)
case nbPoolSizeMedium:
b := (*[nbPoolSizeMedium]byte)(b[0:nbPoolSizeMedium])
nbPoolMedium.Put(b)
case nbPoolSizeLarge:
b := (*[nbPoolSizeLarge]byte)(b[0:nbPoolSizeLarge])
nbPoolLarge.Put(b)
default:
// Ignore frames that are the wrong size, this might happen
// with WebSocket/MQTT messages as they are framed
func nbPoolPut(in []byte) {
ca := cap(in)
for in = in[:ca]; ca >= nbPoolSizeSmall; ca = cap(in) {
switch {
case ca >= nbPoolSizeLarge:
b := (*[nbPoolSizeLarge]byte)(in[0:nbPoolSizeLarge:nbPoolSizeLarge])
nbPoolLarge.Put(b)
in = in[nbPoolSizeLarge:]
case ca >= nbPoolSizeMedium:
b := (*[nbPoolSizeMedium]byte)(in[0:nbPoolSizeMedium:nbPoolSizeMedium])
nbPoolMedium.Put(b)
in = in[nbPoolSizeMedium:]
case ca >= nbPoolSizeSmall:
b := (*[nbPoolSizeSmall]byte)(in[0:nbPoolSizeSmall:nbPoolSizeSmall])
nbPoolSmall.Put(b)
in = in[nbPoolSizeSmall:]
}
}
}

33
server/websocket.go
View File

@@ -115,6 +115,7 @@ type websocket struct {
nocompfrag bool // No fragment for compressed frames
maskread bool
maskwrite bool
compressor *flate.Writer
cookieJwt string
clientIP string
}
@@ -1295,7 +1296,13 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
mfs = 0
}
buf := bytes.NewBuffer(nbPoolGet(usz))
cp, _ := flate.NewWriter(buf, flate.BestSpeed)
cp := c.ws.compressor
if cp == nil {
c.ws.compressor, _ = flate.NewWriter(buf, flate.BestSpeed)
cp = c.ws.compressor
} else {
cp.Reset(buf)
}
var csz int
for _, b := range nb {
cp.Write(b)
@@ -1323,9 +1330,7 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
if mask {
wsMaskBuf(key, p[:lp])
}
new := nbPoolGet(wsFrameSizeForBrowsers)
lp = copy(new[:wsFrameSizeForBrowsers], p[:lp])
bufs = append(bufs, fh[:n], new[:lp])
bufs = append(bufs, fh[:n], p[:lp])
csz += n + lp
p = p[lp:]
}
@@ -1335,16 +1340,9 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
if mask {
wsMaskBuf(key, p)
}
bufs = append(bufs, h)
for len(p) > 0 {
new := nbPoolGet(len(p))
n := copy(new[:cap(new)], p)
bufs = append(bufs, new[:n])
p = p[n:]
}
bufs = append(bufs, h, p)
csz = len(h) + ol
}
nbPoolPut(b) // No longer needed as we copied from above.
// Add to pb the compressed data size (including headers), but
// remove the original uncompressed data size that was added
// during the queueing.
@@ -1355,14 +1353,15 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
if mfs > 0 {
// We are limiting the frame size.
startFrame := func() int {
bufs = append(bufs, nbPoolGet(wsMaxFrameHeaderSize)[:wsMaxFrameHeaderSize])
bufs = append(bufs, nbPoolGet(wsMaxFrameHeaderSize))
return len(bufs) - 1
}
endFrame := func(idx, size int) {
bufs[idx] = bufs[idx][:wsMaxFrameHeaderSize]
n, key := wsFillFrameHeader(bufs[idx], mask, wsFirstFrame, wsFinalFrame, wsUncompressedFrame, wsBinaryMessage, size)
bufs[idx] = bufs[idx][:n]
c.out.pb += int64(n)
c.ws.fs += int64(n + size)
bufs[idx] = bufs[idx][:n]
if mask {
wsMaskBufs(key, bufs[idx+1:])
}
@@ -1388,10 +1387,8 @@ func (c *client) wsCollapsePtoNB() (net.Buffers, int64) {
if endStart {
fhIdx = startFrame()
}
new := nbPoolGet(total)
n := copy(new[:cap(new)], b[:total])
bufs = append(bufs, new[:n])
b = b[n:]
bufs = append(bufs, b[:total])
b = b[total:]
}
}
if total > 0 {