Reuse buffers on socket pumping

proxy/copy_pool.go

@@ -0,0 +1,16 @@
+package proxy
+
+import "sync"
+
+const copyBufferSize = 30 * 1024
+
+var copyPool sync.Pool
+
+func init() {
+	copyPool = sync.Pool{
+		New: func() interface{} {
+			data := make([]byte, copyBufferSize)
+			return &data
+		},
+	}
+}

proxy/server.go

@@ -83,14 +83,10 @@ func (s *Server) accept(conn net.Conn) {
 
 	wait := &sync.WaitGroup{}
 	wait.Add(2)
-	go func() {
-		defer wait.Done()
-		io.Copy(clientConn, tgConn) // nolint: errcheck
-	}()
-	go func() {
-		defer wait.Done()
-		io.Copy(tgConn, clientConn) // nolint: errcheck
-	}()
+
+	go s.pipe(clientConn, tgConn, wait)
+	go s.pipe(tgConn, clientConn, wait)
+
 	<-ctx.Done()
 	wait.Wait()
 
@@ -131,6 +127,15 @@ func (s *Server) getTelegramStream(ctx context.Context, cancel context.CancelFun
 	return conn, nil
 }
 
+func (s *Server) pipe(dst io.Writer, src io.Reader, wait *sync.WaitGroup) {
+	defer wait.Done()
+
+	buf := copyPool.Get().(*[]byte)
+	defer copyPool.Put(buf)
+
+	io.CopyBuffer(dst, src, *buf) // nolint: errcheck
+}
+
 // NewServer creates new instance of MTPROTO proxy.
 func NewServer(conf *config.Config, logger *zap.SugaredLogger, stat *Stats) *Server {
 	return &Server{