cmd serve: pooling copy buffer

rkonfj · May 17, 2023 · e75ee0e · e75ee0e
1 parent ff73a91
commit e75ee0e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 3 deletions.
diff --git a/cmd/serve/cmd.go b/cmd/serve/cmd.go
@@ -1,6 +1,7 @@
 package serve
 
 import (
+	"github.com/dustin/go-humanize"
 	"github.com/rkonfj/toh/server"
 	"github.com/spf13/cobra"
 )
@@ -15,6 +16,7 @@ func init() {
 		RunE:  startAction,
 	}
 	Cmd.Flags().String("acl", "acl.json", "file containing access control rules")
+	Cmd.Flags().String("copy-buf", "16Ki", "buffer size for copying network data")
 	Cmd.Flags().StringP("listen", "l", "localhost:9986", "http server listen address")
 }
 
@@ -37,5 +39,13 @@ func processServerOptions(cmd *cobra.Command) (options server.Options, err error
 		return
 	}
 	options.ACL, err = cmd.Flags().GetString("acl")
+	if err != nil {
+		return
+	}
+	copyBuf, err := cmd.Flags().GetString("copy-buf")
+	if err != nil {
+		return
+	}
+	options.Buf, err = humanize.ParseBytes(copyBuf)
 	return
 }
diff --git a/server/server.go b/server/server.go
@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"io"
+	"math"
 	"net"
 	"net/http"
 	"sync"
@@ -16,11 +17,13 @@ type TohServer struct {
 	options          Options
 	acl              *ACL
 	trafficEventChan chan *TrafficEvent
+	bufPool          *sync.Pool
 }
 
 type Options struct {
 	Listen string
 	ACL    string
+	Buf    uint64
 }
 
 func NewTohServer(options Options) (*TohServer, error) {
@@ -31,7 +34,11 @@ func NewTohServer(options Options) (*TohServer, error) {
 	return &TohServer{
 		options:          options,
 		acl:              acl,
-		trafficEventChan: make(chan *TrafficEvent, 4096),
+		trafficEventChan: make(chan *TrafficEvent, 2048),
+		bufPool: &sync.Pool{New: func() any {
+			buf := make([]byte, int(math.Max(float64(options.Buf), 512)))
+			return &buf
+		}},
 	}, nil
 }
 
@@ -91,11 +98,15 @@ func (s *TohServer) pipe(wsConn *websocket.Conn, netConn net.Conn) (lbc, rbc int
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		lbc, _ = io.Copy(netConn, RWWS(wsConn))
+		buf := s.bufPool.Get().(*[]byte)
+		defer s.bufPool.Put(buf)
+		lbc, _ = io.CopyBuffer(netConn, RWWS(wsConn), *buf)
 		logrus.Debugf("ws conn closed, close remote conn(%s) now", netConn.RemoteAddr().String())
 		netConn.Close()
 	}()
-	rbc, _ = io.Copy(RWWS(wsConn), netConn)
+	buf := s.bufPool.Get().(*[]byte)
+	defer s.bufPool.Put(buf)
+	rbc, _ = io.CopyBuffer(RWWS(wsConn), netConn, *buf)
 	logrus.Debugf("remote conn(%s) closed, close ws conn now", netConn.RemoteAddr().String())
 	wsConn.Close(websocket.StatusBadGateway, "remote close")
 	wg.Wait()