Merge pull request #414 from dolonet/optimize-per-connection-overhead

Reduce per-connection memory overhead

mtglib/internal/doppel/conn.go

@@ -9,6 +9,13 @@ import (
 	"github.com/9seconds/mtg/v2/mtglib/internal/tls"
 )
 
+var doppelBufPool = sync.Pool{
+	New: func() any {
+		b := make([]byte, tls.MaxRecordSize)
+		return &b
+	},
+}
+
 type Conn struct {
 	essentials.Conn
 
@@ -46,7 +53,9 @@ func (c Conn) Start() {
 }
 
 func (c Conn) start() {
-	buf := [tls.MaxRecordSize]byte{}
+	bp := doppelBufPool.Get().(*[]byte)
+	buf := *bp
+	defer doppelBufPool.Put(bp)
 
 	for {
 		select {
@@ -68,7 +77,7 @@ func (c Conn) start() {
 			continue
 		}
 
-		if err := tls.WriteRecordInPlace(c.Conn, buf[:], n); err != nil {
+		if err := tls.WriteRecordInPlace(c.Conn, buf, n); err != nil {
 			c.p.ctxCancel(err)
 			return
 		}

mtglib/internal/relay/relay.go

@@ -4,11 +4,19 @@ import (
 	"context"
 	"errors"
 	"io"
+	"sync"
 
 	"github.com/9seconds/mtg/v2/essentials"
 	"github.com/9seconds/mtg/v2/mtglib/internal/tls"
 )
 
+var bufPool = sync.Pool{
+	New: func() any {
+		b := make([]byte, tls.MaxRecordPayloadSize)
+		return &b
+	},
+}
+
 func Relay(ctx context.Context, log Logger, telegramConn, clientConn essentials.Conn) {
 	defer telegramConn.Close() //nolint: errcheck
 	defer clientConn.Close()   //nolint: errcheck
@@ -16,11 +24,11 @@ func Relay(ctx context.Context, log Logger, telegramConn, clientConn essentials.
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	go func() {
-		<-ctx.Done()
+	stop := context.AfterFunc(ctx, func() {
 		telegramConn.Close() //nolint: errcheck
 		clientConn.Close()   //nolint: errcheck
-	}()
+	})
+	defer stop()
 
 	closeChan := make(chan struct{})
 
@@ -36,12 +44,13 @@ func Relay(ctx context.Context, log Logger, telegramConn, clientConn essentials.
 }
 
 func pump(log Logger, src, dst essentials.Conn, direction string) {
-	var buf [tls.MaxRecordPayloadSize]byte
+	bp := bufPool.Get().(*[]byte)
+	defer bufPool.Put(bp)
 
 	defer src.CloseRead()  //nolint: errcheck
 	defer dst.CloseWrite() //nolint: errcheck
 
-	n, err := io.CopyBuffer(src, dst, buf[:])
+	n, err := io.CopyBuffer(src, dst, *bp)
 
 	switch {
 	case err == nil:

mtglib/proxy.go

@@ -65,10 +65,10 @@ func (p *Proxy) ServeConn(conn essentials.Conn) {
 	ctx := newStreamContext(p.ctx, p.logger, conn)
 	defer ctx.Close()
 
-	go func() {
-		<-ctx.Done()
+	stop := context.AfterFunc(ctx, func() {
 		ctx.Close()
-	}()
+	})
+	defer stop()
 
 	p.eventStream.Send(ctx, NewEventStart(ctx.streamID, ctx.ClientIP()))
 	ctx.logger.Info("Stream has been started")