From ca9241595ebb89ed38c1cef150012a9240b5b352 Mon Sep 17 00:00:00 2001 From: tursom Date: Tue, 23 Jun 2026 12:17:13 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E7=BD=91=E5=85=B3=E8=BD=AC?= =?UTF-8?q?=E5=8F=91=E6=80=A7=E8=83=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cmd/gateway/haproxy.go | 5 +- cmd/gateway/kcp.go | 16 +- cmd/gateway/main.go | 54 +--- cmd/gateway/quic.go | 18 +- cmd/gateway/relay.go | 127 +++++++++ cmd/gateway/relay_benchmark_test.go | 427 ++++++++++++++++++++++++++++ cmd/gateway/tcp.go | 7 +- cmd/gateway/websocket.go | 64 +++-- 8 files changed, 645 insertions(+), 73 deletions(-) create mode 100644 cmd/gateway/relay.go create mode 100644 cmd/gateway/relay_benchmark_test.go diff --git a/cmd/gateway/haproxy.go b/cmd/gateway/haproxy.go index bf54cc4..e21c8d5 100644 --- a/cmd/gateway/haproxy.go +++ b/cmd/gateway/haproxy.go @@ -14,11 +14,12 @@ func haProxyUpstream(source net.Conn, host string) net.Conn { return nil } - conn, err := net.DialTCP("tcp", nil, target) + conn, err := tcpDialer.Dial("tcp", target.String()) if err != nil { log.Err(err).Msg("failed to dial TCP") return nil } + setSocketOptions(conn) sourceAddr, err := net.ResolveTCPAddr( source.RemoteAddr().Network(), @@ -26,6 +27,7 @@ func haProxyUpstream(source net.Conn, host string) net.Conn { ) if err != nil { log.Err(err).Msg("failed to resolve TCP address") + conn.Close() return nil } @@ -45,6 +47,7 @@ func haProxyUpstream(source net.Conn, host string) net.Conn { _, err = header.WriteTo(conn) if err != nil { log.Err(err).Msg("failed to write proxy header") + conn.Close() return nil } diff --git a/cmd/gateway/kcp.go b/cmd/gateway/kcp.go index 33defa5..071991d 100644 --- a/cmd/gateway/kcp.go +++ b/cmd/gateway/kcp.go @@ -31,11 +31,11 @@ func runKcp(wg *sync.WaitGroup) { Msg("Failed to accept KCP connection") continue } - log.Info(). + log.Debug(). Str("remote_addr", conn.RemoteAddr().String()). Msg("Accepted KCP connection") - conn.SetACKNoDelay(true) + tuneKcpConn(conn) go handleRequest(conn) } @@ -46,9 +46,17 @@ func upstreamKcp(host string) net.Conn { if err != nil { log.Error().Err(err). Msg("Failed to dial KCP server") + return nil } - defer conn.Close() - conn.SetACKNoDelay(true) + tuneKcpConn(conn) return conn } + +func tuneKcpConn(conn *kcp.UDPSession) { + conn.SetStreamMode(true) + conn.SetWriteDelay(false) + conn.SetNoDelay(1, 10, 2, 1) + conn.SetWindowSize(256, 256) + conn.SetACKNoDelay(true) +} diff --git a/cmd/gateway/main.go b/cmd/gateway/main.go index afb1bd9..77d96d8 100644 --- a/cmd/gateway/main.go +++ b/cmd/gateway/main.go @@ -1,12 +1,9 @@ package main import ( - "io" "net" "strings" - "sync" - "github.com/rs/zerolog" "github.com/rs/zerolog/log" "github.com/tursom/mc-gateway/plugin/api" "github.com/tursom/mc-gateway/protocol" @@ -66,19 +63,13 @@ func handleRequest(conn net.Conn) { } defer client.Close() - var wg sync.WaitGroup - - wg.Add(1) - go copyData(client, conn, &wg) - copyData(conn, client, nil) - - // 等待所有读写操作完成 - // 不放在 defer 中，以防报错时无法关闭连接 - wg.Wait() + proxyConnections(conn, client) } func mapToHost(conn net.Conn) net.Conn { - buf := make([]byte, 1024) + buf := getProxyBuffer() + defer putProxyBuffer(buf) + n, err := conn.Read(buf) if err != nil { log.Err(err). @@ -113,7 +104,7 @@ func mapToHost(conn net.Conn) net.Conn { return nil } - log.Info(). + log.Debug(). Str("client", conn.RemoteAddr().String()). Str("host", mc_host). Str("mc", host). @@ -146,32 +137,15 @@ func mapToHost(conn net.Conn) net.Conn { return nil } - client.Write(buf[:n]) + if err := writeAll(client, buf[:n]); err != nil { + log.Err(err). + Str("client", conn.RemoteAddr().String()). + Str("host", mc_host). + Str("mc", host). + Msg("failed to write initial packet to upstream") + client.Close() + return nil + } return client } - -func copyData(dst io.Writer, src io.Reader, wg *sync.WaitGroup) { - defer func() { - if r := recover(); r != nil { - var event *zerolog.Event - if err, ok := r.(error); ok { - event = log.Err(err) - } else if str, ok := r.(string); ok { - event = log.Error().Str("panic", str) - } else { - event = log.Error().Any("panic", r) - } - event.Msg("Panic in copyData") - } - }() - - if wg != nil { - defer wg.Done() - } - - _, err := io.Copy(dst, src) - if err != nil && err != io.EOF { - log.Err(err).Msg("Error copying data") - } -} diff --git a/cmd/gateway/quic.go b/cmd/gateway/quic.go index aa6f5e2..9720d22 100644 --- a/cmd/gateway/quic.go +++ b/cmd/gateway/quic.go @@ -54,7 +54,7 @@ func runQuic(wg *sync.WaitGroup) { continue } - log.Info(). + log.Debug(). Str("client", conn.RemoteAddr().String()). Msg("Accepted QUIC connection") go handleQuicRequest(conn) @@ -79,9 +79,10 @@ func upstreamQuic(host string) net.Conn { stream, err := conn.OpenStream() if err != nil { log.Err(err).Str("host", host).Msg("Failed to open stream") + conn.CloseWithError(0, "failed to open stream") return nil } - log.Info().Str("host", host).Msg("QUIC stream opened") + log.Debug().Str("host", host).Msg("QUIC stream opened") return quicConn{ Connection: conn, @@ -164,8 +165,15 @@ func getQuicNextProtos() []string { } func (c quicConn) Close() error { - if err := c.Stream.Close(); err != nil { - log.Err(err).Msg("Failed to close QUIC stream") - } + _ = c.Stream.Close() return c.Connection.CloseWithError(0, "Closing QUIC connection") } + +func (c quicConn) CloseWrite() error { + return c.Stream.Close() +} + +func (c quicConn) CloseRead() error { + c.Stream.CancelRead(0) + return nil +} diff --git a/cmd/gateway/relay.go b/cmd/gateway/relay.go new file mode 100644 index 0000000..efdf6b8 --- /dev/null +++ b/cmd/gateway/relay.go @@ -0,0 +1,127 @@ +package main + +import ( + "errors" + "io" + "sync" + + "github.com/rs/zerolog/log" +) + +const proxyBufferSize = 64 * 1024 + +var proxyBufferPool = sync.Pool{ + New: func() any { + buf := make([]byte, proxyBufferSize) + return &buf + }, +} + +type ( + closeWriter interface { + CloseWrite() error + } + + closeReader interface { + CloseRead() error + } +) + +func proxyConnections(a, b io.ReadWriter) { + var wg sync.WaitGroup + + wg.Add(2) + go func() { + defer wg.Done() + proxyCopy(b, a) + }() + go func() { + defer wg.Done() + proxyCopy(a, b) + }() + + wg.Wait() +} + +func proxyCopy(dst io.Writer, src io.Reader) { + defer recoverProxyCopy() + defer closeRead(src) + defer closeWrite(dst) + + _, err := copyForward(dst, src) + if err != nil && !errors.Is(err, io.EOF) { + log.Debug().Err(err).Msg("proxy copy stopped") + } +} + +func copyForward(dst io.Writer, src io.Reader) (int64, error) { + if _, ok := src.(io.WriterTo); ok { + return io.Copy(dst, src) + } + if _, ok := dst.(io.ReaderFrom); ok { + return io.Copy(dst, src) + } + + buf := getProxyBuffer() + defer putProxyBuffer(buf) + + return io.CopyBuffer(dst, src, buf) +} + +func getProxyBuffer() []byte { + return *proxyBufferPool.Get().(*[]byte) +} + +func putProxyBuffer(buf []byte) { + if cap(buf) != proxyBufferSize { + return + } + buf = buf[:proxyBufferSize] + proxyBufferPool.Put(&buf) +} + +func writeAll(w io.Writer, buf []byte) error { + for len(buf) > 0 { + n, err := w.Write(buf) + if n > 0 { + buf = buf[n:] + } + if err != nil { + return err + } + if n == 0 { + return io.ErrShortWrite + } + } + + return nil +} + +func closeWrite(conn any) { + if closer, ok := conn.(closeWriter); ok { + if err := closer.CloseWrite(); err != nil { + log.Debug().Err(err).Msg("failed to close write side") + } + return + } + + if closer, ok := conn.(io.Closer); ok { + if err := closer.Close(); err != nil { + log.Debug().Err(err).Msg("failed to close connection") + } + } +} + +func closeRead(conn any) { + if closer, ok := conn.(closeReader); ok { + if err := closer.CloseRead(); err != nil { + log.Debug().Err(err).Msg("failed to close read side") + } + } +} + +func recoverProxyCopy() { + if rec := recover(); rec != nil { + log.Error().Any("panic", rec).Msg("panic in proxy copy") + } +} diff --git a/cmd/gateway/relay_benchmark_test.go b/cmd/gateway/relay_benchmark_test.go new file mode 100644 index 0000000..34a765e --- /dev/null +++ b/cmd/gateway/relay_benchmark_test.go @@ -0,0 +1,427 @@ +package main + +import ( + "bytes" + "errors" + "io" + "net" + "testing" + "time" + + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" + "github.com/tursom/mc-gateway/plugin/api" +) + +const benchmarkPayloadSize = 4 << 20 + +func BenchmarkForwardCopy(b *testing.B) { + disableBenchmarkLogs(b) + + b.Run("stdlib_io_copy", func(b *testing.B) { + benchmarkCopy(b, func(dst io.Writer, src io.Reader) { + if _, err := io.Copy(dst, src); err != nil { + b.Fatal(err) + } + }) + }) + + b.Run("pooled_copy_buffer", func(b *testing.B) { + benchmarkCopy(b, func(dst io.Writer, src io.Reader) { + proxyCopy(dst, src) + }) + }) +} + +func BenchmarkTCPForwardCopy(b *testing.B) { + disableBenchmarkLogs(b) + + b.Run("generic_user_buffer", func(b *testing.B) { + benchmarkTCPForward(b, func(dst, src *net.TCPConn) error { + defer closeRead(src) + defer closeWrite(dst) + + buf := getProxyBuffer() + defer putProxyBuffer(buf) + + _, err := io.CopyBuffer( + benchmarkTCPWriterOnly{conn: dst}, + benchmarkTCPReaderOnly{conn: src}, + buf, + ) + return err + }) + }) + + b.Run("tcp_fast_path", func(b *testing.B) { + benchmarkTCPForward(b, func(dst, src *net.TCPConn) error { + proxyCopy(dst, src) + return nil + }) + }) +} + +func BenchmarkMapToHostInitialPacket(b *testing.B) { + disableBenchmarkLogs(b) + + const ( + hostName = "dev.example" + upstreamHost = "benchmark-upstream" + ) + + packet := benchmarkHandshakePacket(hostName) + source := newBenchmarkConn(benchmarkAddr("client:25565")) + upstream := newBenchmarkConn(benchmarkAddr("upstream:25565")) + + restore := installBenchmarkUpstreamHook(b, hostName, upstreamHost, upstream) + defer restore() + + b.SetBytes(int64(len(packet))) + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + source.ResetReader(packet) + upstream.ResetWriter() + + client := mapToHost(source) + if client == nil { + b.Fatal("mapToHost returned nil") + } + if upstream.Written() != int64(len(packet)) { + b.Fatalf("upstream wrote %d bytes, want %d", upstream.Written(), len(packet)) + } + } +} + +func benchmarkCopy(b *testing.B, copyFunc func(io.Writer, io.Reader)) { + reader := newBenchmarkReader(benchmarkPayloadSize) + writer := &benchmarkWriter{} + + b.SetBytes(benchmarkPayloadSize) + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + reader.Reset(benchmarkPayloadSize) + writer.Reset() + + copyFunc(writer, reader) + if writer.Written() != benchmarkPayloadSize { + b.Fatalf("copied %d bytes, want %d", writer.Written(), benchmarkPayloadSize) + } + } +} + +func benchmarkTCPForward(b *testing.B, copyFunc func(dst, src *net.TCPConn) error) { + b.SetBytes(benchmarkPayloadSize) + b.ReportAllocs() + b.ResetTimer() + + for i := 0; i < b.N; i++ { + b.StopTimer() + + srcRelay, srcWriter := newBenchmarkTCPConnPair(b) + dstReader, dstRelay := newBenchmarkTCPConnPair(b) + setBenchmarkTCPDeadline(b, srcRelay, srcWriter, dstReader, dstRelay) + + start := make(chan struct{}) + writerDone := make(chan error, 1) + readerDone := make(chan benchmarkTCPReadResult, 1) + + go benchmarkTCPWrite(start, srcWriter, writerDone) + go benchmarkTCPRead(start, dstReader, readerDone) + + b.StartTimer() + close(start) + + copyErr := copyFunc(dstRelay, srcRelay) + writeErr := <-writerDone + readResult := <-readerDone + + b.StopTimer() + closeBenchmarkTCPConns(srcRelay, srcWriter, dstReader, dstRelay) + + if copyErr != nil && !errors.Is(copyErr, io.EOF) { + b.Fatalf("copy failed: %v", copyErr) + } + if writeErr != nil { + b.Fatalf("source write failed: %v", writeErr) + } + if readResult.err != nil { + b.Fatalf("destination read failed: %v", readResult.err) + } + if readResult.n != benchmarkPayloadSize { + b.Fatalf("forwarded %d bytes, want %d", readResult.n, benchmarkPayloadSize) + } + } +} + +func benchmarkTCPWrite(start <-chan struct{}, conn *net.TCPConn, done chan<- error) { + <-start + + reader := newBenchmarkReader(benchmarkPayloadSize) + buf := getProxyBuffer() + _, err := io.CopyBuffer(conn, reader, buf) + putProxyBuffer(buf) + + if closeErr := conn.CloseWrite(); err == nil { + err = closeErr + } + done <- err +} + +func benchmarkTCPRead(start <-chan struct{}, conn *net.TCPConn, done chan<- benchmarkTCPReadResult) { + <-start + + n, err := io.Copy(io.Discard, conn) + done <- benchmarkTCPReadResult{n: n, err: err} +} + +func benchmarkHandshakePacket(host string) []byte { + packet := make([]byte, 5+len(host)+2) + packet[4] = byte(len(host)) + copy(packet[5:], host) + return packet +} + +func installBenchmarkUpstreamHook(b *testing.B, hostName, upstreamHost string, upstream net.Conn) func() { + b.Helper() + + previousConfig := config + previousHooks := hooks + + config.Hosts = map[string]string{ + hostName: upstreamHost, + } + + pluginLock.Lock() + hooks = map[string]map[string]any{ + "benchmark": {}, + } + pluginLock.Unlock() + + gateway := &Gateway{pluginId: "benchmark"} + if err := api.RegisterHookHandler( + gateway, + api.HookUpstream, + func(_ net.Conn, host string) bool { + return host == upstreamHost + }, + func(_ net.Conn, _ string) (net.Conn, error) { + return upstream, nil + }, + ); err != nil { + b.Fatal(err) + } + + return func() { + config = previousConfig + + pluginLock.Lock() + hooks = previousHooks + pluginLock.Unlock() + } +} + +func newBenchmarkTCPConnPair(b *testing.B) (*net.TCPConn, *net.TCPConn) { + b.Helper() + + listener, err := net.ListenTCP("tcp4", &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)}) + if err != nil { + b.Fatal(err) + } + defer listener.Close() + + accepted := make(chan *net.TCPConn, 1) + acceptErr := make(chan error, 1) + go func() { + conn, err := listener.AcceptTCP() + if err != nil { + acceptErr <- err + return + } + accepted <- conn + }() + + client, err := net.DialTCP("tcp4", nil, listener.Addr().(*net.TCPAddr)) + if err != nil { + b.Fatal(err) + } + + select { + case err := <-acceptErr: + client.Close() + b.Fatal(err) + case server := <-accepted: + setSocketOptions(server) + setSocketOptions(client) + return server, client + } + + return nil, nil +} + +func setBenchmarkTCPDeadline(b *testing.B, conns ...*net.TCPConn) { + b.Helper() + + deadline := time.Now().Add(30 * time.Second) + for _, conn := range conns { + if err := conn.SetDeadline(deadline); err != nil { + b.Fatal(err) + } + } +} + +func closeBenchmarkTCPConns(conns ...*net.TCPConn) { + for _, conn := range conns { + conn.Close() + } +} + +func disableBenchmarkLogs(b *testing.B) { + b.Helper() + + previousLogger := log.Logger + log.Logger = zerolog.Nop() + b.Cleanup(func() { + log.Logger = previousLogger + }) +} + +type benchmarkTCPReadResult struct { + n int64 + err error +} + +type benchmarkTCPReaderOnly struct { + conn *net.TCPConn +} + +func (r benchmarkTCPReaderOnly) Read(p []byte) (int, error) { + return r.conn.Read(p) +} + +type benchmarkTCPWriterOnly struct { + conn *net.TCPConn +} + +func (w benchmarkTCPWriterOnly) Write(p []byte) (int, error) { + return w.conn.Write(p) +} + +type benchmarkReader struct { + chunk []byte + remaining int64 +} + +func newBenchmarkReader(size int64) *benchmarkReader { + chunk := bytes.Repeat([]byte{0x5a}, 64*1024) + return &benchmarkReader{ + chunk: chunk, + remaining: size, + } +} + +func (r *benchmarkReader) Reset(size int64) { + r.remaining = size +} + +func (r *benchmarkReader) Read(p []byte) (int, error) { + if r.remaining == 0 { + return 0, io.EOF + } + + n := len(p) + if n > len(r.chunk) { + n = len(r.chunk) + } + if int64(n) > r.remaining { + n = int(r.remaining) + } + + copy(p, r.chunk[:n]) + r.remaining -= int64(n) + return n, nil +} + +type benchmarkWriter struct { + written int64 +} + +func (w *benchmarkWriter) Reset() { + w.written = 0 +} + +func (w *benchmarkWriter) Written() int64 { + return w.written +} + +func (w *benchmarkWriter) Write(p []byte) (int, error) { + w.written += int64(len(p)) + return len(p), nil +} + +type benchmarkConn struct { + reader bytes.Reader + writer benchmarkWriter + addr net.Addr +} + +func newBenchmarkConn(addr net.Addr) *benchmarkConn { + return &benchmarkConn{addr: addr} +} + +func (c *benchmarkConn) ResetReader(buf []byte) { + c.reader.Reset(buf) +} + +func (c *benchmarkConn) ResetWriter() { + c.writer.Reset() +} + +func (c *benchmarkConn) Written() int64 { + return c.writer.Written() +} + +func (c *benchmarkConn) Read(p []byte) (int, error) { + return c.reader.Read(p) +} + +func (c *benchmarkConn) Write(p []byte) (int, error) { + return c.writer.Write(p) +} + +func (c *benchmarkConn) Close() error { + return nil +} + +func (c *benchmarkConn) LocalAddr() net.Addr { + return c.addr +} + +func (c *benchmarkConn) RemoteAddr() net.Addr { + return c.addr +} + +func (c *benchmarkConn) SetDeadline(time.Time) error { + return nil +} + +func (c *benchmarkConn) SetReadDeadline(time.Time) error { + return nil +} + +func (c *benchmarkConn) SetWriteDeadline(time.Time) error { + return nil +} + +type benchmarkAddr string + +func (a benchmarkAddr) Network() string { + return "tcp" +} + +func (a benchmarkAddr) String() string { + return string(a) +} diff --git a/cmd/gateway/tcp.go b/cmd/gateway/tcp.go index bef9902..2be2f49 100644 --- a/cmd/gateway/tcp.go +++ b/cmd/gateway/tcp.go @@ -39,7 +39,7 @@ func runTcp(wg *sync.WaitGroup) { } func upstreamTcp(host string) net.Conn { - conn, err := net.Dial("tcp", host) + conn, err := tcpDialer.Dial("tcp", host) if err != nil { log.Err(err).Str("host", host).Msg("Error dialing upstream") return nil @@ -49,6 +49,11 @@ func upstreamTcp(host string) net.Conn { } +var tcpDialer = net.Dialer{ + Timeout: 3 * time.Second, + KeepAlive: 30 * time.Second, +} + func setSocketOptions(conn net.Conn) { if tcpConn, ok := conn.(*net.TCPConn); ok { tcpConn.SetNoDelay(true) // 禁用 Nagle 算法 diff --git a/cmd/gateway/websocket.go b/cmd/gateway/websocket.go index 5c31344..117f2e0 100644 --- a/cmd/gateway/websocket.go +++ b/cmd/gateway/websocket.go @@ -1,7 +1,9 @@ package main import ( + "errors" "fmt" + "io" "net/http" "sync" "time" @@ -21,7 +23,7 @@ type ( // WebSocket 连接适配器，实现 net.Conn 接口 webSocketConn struct { *websocket.Conn - messageRemain []byte // 用于存储未处理的消息 + reader io.Reader } ) @@ -39,38 +41,56 @@ func handleWebSocket(w http.ResponseWriter, r *http.Request) { } func (w *webSocketConn) Read(b []byte) (n int, err error) { - // 如果有未处理的消息，直接从 messageRemain 中读取 - if len(w.messageRemain) > 0 { - copied := copy(b, w.messageRemain) - if copied < len(w.messageRemain) { - w.messageRemain = w.messageRemain[copied:] - } else { - w.messageRemain = nil // 清空已处理的消息 + for { + if w.reader != nil { + n, err = w.reader.Read(b) + if errors.Is(err, io.EOF) { + w.reader = nil + if n > 0 { + return n, nil + } + continue + } + return n, err } - return copied, nil - } - // 读取新消息 - _, message, err := w.ReadMessage() - if err != nil { - return 0, err + messageType, reader, err := w.NextReader() + if err != nil { + return 0, err + } + if messageType != websocket.BinaryMessage && messageType != websocket.TextMessage { + continue + } + w.reader = reader } - copied := copy(b, message) - if copied < len(message) { - w.messageRemain = message[copied:] - } - return copied, nil } func (w *webSocketConn) Write(b []byte) (n int, err error) { - if err = w.WriteMessage(websocket.BinaryMessage, b); err != nil { + writer, err := w.NextWriter(websocket.BinaryMessage) + if err != nil { return 0, err } - return len(b), nil + + n, err = writer.Write(b) + closeErr := writer.Close() + if err != nil { + return n, err + } + if closeErr != nil { + return n, closeErr + } + if n != len(b) { + return n, io.ErrShortWrite + } + + return n, nil } func (w *webSocketConn) SetDeadline(t time.Time) error { - return w.SetReadDeadline(t) + if err := w.SetReadDeadline(t); err != nil { + return err + } + return w.SetWriteDeadline(t) } // 启动 WebSocket 服务器