diff --git a/mtglib/internal/doppel/conn.go b/mtglib/internal/doppel/conn.go index cc81d91..e4342ba 100644 --- a/mtglib/internal/doppel/conn.go +++ b/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 } diff --git a/mtglib/internal/relay/relay.go b/mtglib/internal/relay/relay.go index 986b3e2..70ee140 100644 --- a/mtglib/internal/relay/relay.go +++ b/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: diff --git a/mtglib/proxy.go b/mtglib/proxy.go index be05194..44d426f 100644 --- a/mtglib/proxy.go +++ b/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")