diff --git a/rpc/server.go b/rpc/server.go index c2245af6b7..1d5b1663f2 100644 --- a/rpc/server.go +++ b/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. diff --git a/rpc/websocket.go b/rpc/websocket.go index 15bfb18495..0fb203e6eb 100644 --- a/rpc/websocket.go +++ b/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, - WriteBufferSize: wsWriteBuffer, - WriteBufferPool: wsBufferPool, - CheckOrigin: wsHandshakeValidator(allowedOrigins), + EnableCompression: enableCompression, + ReadBufferSize: wsReadBuffer, + WriteBufferSize: wsWriteBuffer, + 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) })