feat(rpc): allow configure WS compression (#210)

* Default open ws compression.

* Default open ws compression.

* Enable ws compression.

* Enable ws compression.

---------

Co-authored-by: HAOYUatHZ <37070449+HAOYUatHZ@users.noreply.github.com>
Co-authored-by: HAOYUatHZ <haoyu@protonmail.com>

rpc/server.go

@@ -17,7 +17,9 @@
 package rpc
 
 import (
+	"compress/flate"
 	"context"
+	"errors"
 	"io"
 	"sync/atomic"
 
@@ -47,6 +49,8 @@ type Server struct {
 	idgen    func() ID
 	run      int32
 	codecs   mapset.Set
+	// Add compressionLevel inorder to enable set it when open websocket server.
+	compressionLevel int
 }
 
 // NewServer creates a new server instance with no registered handlers.
@@ -89,6 +93,15 @@ func (s *Server) ServeCodec(codec ServerCodec, options CodecOption) {
 	c.Close()
 }
 
+// SetCompressionLevel set compression level (-2 ~ 9), this function only works on websocket.
+func (s *Server) SetCompressionLevel(level int) error {
+	if !(flate.HuffmanOnly <= level && level <= flate.BestCompression) {
+		return errors.New("websocket: invalid compression level")
+	}
+	s.compressionLevel = level
+	return nil
+}
+
 // serveSingleRequest reads and processes a single RPC request from the given codec. This
 // is used to serve HTTP connections. Subscriptions and reverse calls are not allowed in
 // this mode.

rpc/websocket.go

@@ -17,6 +17,7 @@
 package rpc
 
 import (
+	"compress/flate"
 	"context"
 	"encoding/base64"
 	"fmt"
@@ -49,11 +50,13 @@ var wsBufferPool = new(sync.Pool)
 // allowedOrigins should be a comma-separated list of allowed origin URLs.
 // To allow connections with any origin, pass "*".
 func (s *Server) WebsocketHandler(allowedOrigins []string) http.Handler {
+	enableCompression := s.compressionLevel != flate.NoCompression
 	var upgrader = websocket.Upgrader{
-		ReadBufferSize:  wsReadBuffer,
+		EnableCompression: enableCompression,
-		WriteBufferSize: wsWriteBuffer,
+		ReadBufferSize:    wsReadBuffer,
-		WriteBufferPool: wsBufferPool,
+		WriteBufferSize:   wsWriteBuffer,
-		CheckOrigin:     wsHandshakeValidator(allowedOrigins),
+		WriteBufferPool:   wsBufferPool,
+		CheckOrigin:       wsHandshakeValidator(allowedOrigins),
 	}
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		conn, err := upgrader.Upgrade(w, r, nil)
@@ -61,6 +64,9 @@ func (s *Server) WebsocketHandler(allowedOrigins []string) http.Handler {
 			log.Debug("WebSocket upgrade failed", "err", err)
 			return
 		}
+		if enableCompression {
+			_ = conn.SetCompressionLevel(s.compressionLevel)
+		}
 		codec := newWebsocketCodec(conn)
 		s.ServeCodec(codec, 0)
 	})