feat: add DialContext and EnrollContext for Client

client_test.go

@@ -7,7 +7,6 @@ import (
 	"io"
 	"math/rand"
 	"net"
-	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -21,12 +20,17 @@ import (
 	goPool "github.com/panjf2000/gnet/v2/pkg/pool/goroutine"
 )
 
+type connHandler struct {
+	network string
+	rspCh   chan []byte
+	data    []byte
+}
+
 type clientEvents struct {
 	*BuiltinEventEngine
 	tester    *testing.T
 	svr       *testClientServer
 	packetLen int
-	rspChMap  sync.Map
 }
 
 func (ev *clientEvents) OnBoot(e Engine) Action {
@@ -38,9 +42,6 @@ func (ev *clientEvents) OnBoot(e Engine) Action {
 }
 
 func (ev *clientEvents) OnOpen(c Conn) ([]byte, Action) {
-	c.SetContext([]byte{})
-	rspCh := make(chan []byte, 1)
-	ev.rspChMap.Store(c.LocalAddr().String(), rspCh)
 	return nil, None
 }
 
@@ -54,24 +55,18 @@ func (ev *clientEvents) OnClose(Conn, error) Action {
 }
 
 func (ev *clientEvents) OnTraffic(c Conn) (action Action) {
-	ctx := c.Context()
-	var p []byte
-	if ctx != nil {
-		p = ctx.([]byte)
-	} else { // UDP
+	handler := c.Context().(*connHandler)
+	if handler.network == "udp" {
 		ev.packetLen = 1024
 	}
 	buf, err := c.Next(-1)
 	assert.NoError(ev.tester, err)
-	p = append(p, buf...)
-	if len(p) < ev.packetLen {
-		c.SetContext(p)
+	handler.data = append(handler.data, buf...)
+	if len(handler.data) < ev.packetLen {
 		return
 	}
-	v, _ := ev.rspChMap.Load(c.LocalAddr().String())
-	rspCh := v.(chan []byte)
-	rspCh <- p
-	c.SetContext([]byte{})
+	handler.rspCh <- handler.data
+	handler.data = handler.data[:0]
 	return
 }
 
@@ -330,38 +325,23 @@ func startGnetClient(t *testing.T, cli *Client, ev *clientEvents, network, addr
 		c   Conn
 		err error
 	)
+	var handler = &connHandler{
+		network: network,
+		rspCh:   make(chan []byte, 1),
+	}
 	if netDial {
 		var netConn net.Conn
 		netConn, err = NetDial(network, addr)
 		require.NoError(t, err)
-		c, err = cli.Enroll(netConn)
+		c, err = cli.EnrollWithContext(netConn, handler)
 	} else {
-		c, err = cli.Dial(network, addr)
+		c, err = cli.DialWithContext(network, addr, handler)
 	}
 	require.NoError(t, err)
 	defer c.Close()
 	err = c.Wake(nil)
 	require.NoError(t, err)
-	var rspCh chan []byte
-	if network == "udp" {
-		rspCh = make(chan []byte, 1)
-		ev.rspChMap.Store(c.LocalAddr().String(), rspCh)
-	} else {
-		var (
-			v  interface{}
-			ok bool
-		)
-		start := time.Now()
-		for time.Since(start) < time.Second {
-			v, ok = ev.rspChMap.Load(c.LocalAddr().String())
-			if ok {
-				break
-			}
-			time.Sleep(10 * time.Millisecond)
-		}
-		require.True(t, ok)
-		rspCh = v.(chan []byte)
-	}
+	var rspCh = handler.rspCh
 	duration := time.Duration((rand.Float64()*2+1)*float64(time.Second)) / 2
 	t.Logf("test duration: %dms", duration/time.Millisecond)
 	start := time.Now()

client_unix.go

@@ -140,15 +140,23 @@ func (cli *Client) Stop() (err error) {
 
 // Dial is like net.Dial().
 func (cli *Client) Dial(network, address string) (Conn, error) {
+	return cli.DialWithContext(network, address, nil)
+}
+
+func (cli *Client) DialWithContext(network, address string, ctx interface{}) (Conn, error) {
 	c, err := net.Dial(network, address)
 	if err != nil {
 		return nil, err
 	}
-	return cli.Enroll(c)
+	return cli.EnrollWithContext(c, ctx)
 }
 
 // Enroll converts a net.Conn to gnet.Conn and then adds it into Client.
 func (cli *Client) Enroll(c net.Conn) (Conn, error) {
+	return cli.EnrollWithContext(c, nil)
+}
+
+func (cli *Client) EnrollWithContext(c net.Conn, ctx interface{}) (Conn, error) {
 	defer c.Close()
 
 	sc, ok := c.(syscall.Conn)
@@ -222,5 +230,6 @@ func (cli *Client) Enroll(c net.Conn) (Conn, error) {
 		gc.Close()
 		return nil, err
 	}
+	gc.SetContext(ctx)
 	return gc, nil
 }

client_windows.go

@@ -118,6 +118,10 @@ func unixAddr(addr string) string {
 }
 
 func (cli *Client) Dial(network, addr string) (Conn, error) {
+	return cli.DialWithContext(network, addr, nil)
+}
+
+func (cli *Client) DialWithContext(network, addr string, ctx interface{}) (Conn, error) {
 	var (
 		c   net.Conn
 		err error
@@ -135,10 +139,14 @@ func (cli *Client) Dial(network, addr string) (Conn, error) {
 			return nil, err
 		}
 	}
-	return cli.Enroll(c)
+	return cli.EnrollWithContext(c, ctx)
 }
 
 func (cli *Client) Enroll(nc net.Conn) (gc Conn, err error) {
+	return cli.EnrollWithContext(nc, nil)
+}
+
+func (cli *Client) EnrollWithContext(nc net.Conn, ctx interface{}) (gc Conn, err error) {
 	switch v := nc.(type) {
 	case *net.TCPConn:
 		if cli.opts.TCPNoDelay == TCPNoDelay {
@@ -156,6 +164,7 @@ func (cli *Client) Enroll(nc net.Conn) (gc Conn, err error) {
 		}
 
 		c := newTCPConn(nc, cli.el)
+		c.SetContext(ctx)
 		cli.el.ch <- c
 		go func(c *conn, tc net.Conn, el *eventloop) {
 			var buffer [0x10000]byte
@@ -171,6 +180,7 @@ func (cli *Client) Enroll(nc net.Conn) (gc Conn, err error) {
 		gc = c
 	case *net.UnixConn:
 		c := newTCPConn(nc, cli.el)
+		c.SetContext(ctx)
 		cli.el.ch <- c
 		go func(c *conn, uc net.Conn, el *eventloop) {
 			var buffer [0x10000]byte
@@ -192,6 +202,7 @@ func (cli *Client) Enroll(nc net.Conn) (gc Conn, err error) {
 		gc = c
 	case *net.UDPConn:
 		c := newUDPConn(cli.el, nc.LocalAddr(), nc.RemoteAddr())
+		c.SetContext(ctx)
 		c.rawConn = nc
 		go func(uc net.Conn, el *eventloop) {
 			var buffer [0x10000]byte
@@ -201,6 +212,7 @@ func (cli *Client) Enroll(nc net.Conn) (gc Conn, err error) {
 					return
 				}
 				c := newUDPConn(cli.el, uc.LocalAddr(), uc.RemoteAddr())
+				c.SetContext(ctx)
 				c.rawConn = uc
 				el.ch <- packUDPConn(c, buffer[:n])
 			}