优化网关转发性能

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
 	}
 

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)
+}

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")
-	}
-}

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
+}

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")
+	}
+}

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)
+}

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 算法

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 服务器