go/neo/neonet: errgroup -> xsync.WorkGroup

See go123@515a6d14

go/neo/neonet/connection_test.go

@@ -29,10 +29,9 @@ import (
 	"testing"
 	"time"
 
-	"golang.org/x/sync/errgroup"
-
 	"lab.nexedi.com/kirr/go123/exc"
 	"lab.nexedi.com/kirr/go123/xerr"
+	"lab.nexedi.com/kirr/go123/xsync"
 
 	"lab.nexedi.com/kirr/neo/go/internal/packed"
 
@@ -76,8 +75,12 @@ func xwait(w interface{ Wait() error }) {
 	exc.Raiseif(err)
 }
 
-func gox(wg interface{ Go(func() error) }, xf func()) {
-	wg.Go(exc.Funcx(xf))
+func gox(wg interface { Go(func(context.Context) error) }, xf func(context.Context)) {
+	wg.Go(func(ctx context.Context) error {
+		return exc.Funcx(func() {
+			xf(ctx)
+		})()
+	})
 }
 
 // xlinkError verifies that err is *LinkError and returns err.Err .
@@ -171,11 +174,12 @@ func nodeLinkPipe() (nl1, nl2 *NodeLink) {
 
 func TestNodeLink(t *testing.T) {
 	// TODO catch exception -> add proper location from it -> t.Fatal (see git-backup)
+	bg := context.Background()
 
 	// Close vs recvPkt
 	nl1, nl2 := _nodeLinkPipe(linkNoRecvSend, linkNoRecvSend)
-	wg := &errgroup.Group{}
-	gox(wg, func() {
+	wg := xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(nl1)
 	})
@@ -188,8 +192,8 @@ func TestNodeLink(t *testing.T) {
 
 	// Close vs sendPkt
 	nl1, nl2 = _nodeLinkPipe(linkNoRecvSend, linkNoRecvSend)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(nl1)
 	})
@@ -203,8 +207,8 @@ func TestNodeLink(t *testing.T) {
 
 	// {Close,CloseAccept} vs Accept
 	nl1, nl2 = _nodeLinkPipe(linkNoRecvSend, linkNoRecvSend)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(nl2)
 	})
@@ -212,7 +216,7 @@ func TestNodeLink(t *testing.T) {
 	if !(c == nil && xlinkError(err) == ErrLinkClosed) {
 		t.Fatalf("NodeLink.Accept() after close: conn = %v, err = %v", c, err)
 	}
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		nl1.CloseAccept()
 	})
@@ -231,8 +235,8 @@ func TestNodeLink(t *testing.T) {
 
 	// Close vs recvPkt on another side
 	nl1, nl2 = _nodeLinkPipe(linkNoRecvSend, linkNoRecvSend)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(nl2)
 	})
@@ -245,8 +249,8 @@ func TestNodeLink(t *testing.T) {
 
 	// Close vs sendPkt on another side
 	nl1, nl2 = _nodeLinkPipe(linkNoRecvSend, linkNoRecvSend)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(nl2)
 	})
@@ -261,40 +265,54 @@ func TestNodeLink(t *testing.T) {
 	// raw exchange
 	nl1, nl2 = _nodeLinkPipe(linkNoRecvSend, linkNoRecvSend)
 
-	wg, ctx := errgroup.WithContext(context.Background())
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	okch := make(chan int, 2)
+	gox(wg, func(_ context.Context) {
 		// send ping; wait for pong
 		pkt := _mkpkt(1, 2, []byte("ping"))
 		xsendPkt(nl1, pkt)
 		pkt = xrecvPkt(nl1)
 		xverifyPkt(pkt, 3, 4, []byte("pong"))
+		okch <- 1
 	})
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		// wait for ping; send pong
 		pkt = xrecvPkt(nl2)
 		xverifyPkt(pkt, 1, 2, []byte("ping"))
 		pkt = _mkpkt(3, 4, []byte("pong"))
 		xsendPkt(nl2, pkt)
+		okch <- 2
 	})
 
 	// close nodelinks either when checks are done, or upon first error
-	wgclose := &errgroup.Group{}
-	gox(wgclose, func() {
-		<-ctx.Done()
+	gox(wg, func(ctx context.Context) {
+		nok := 0
+	loop:
+		for {
+			select {
+			case <-ctx.Done():
+				break loop
+			case <-okch:
+				nok++
+				if nok == 2 {
+					break loop
+				}
+			}
+		}
+
 		xclose(nl1)
 		xclose(nl2)
 	})
 
 	xwait(wg)
-	xwait(wgclose)
 
 	// ---- connections on top of nodelink ----
 
 	// Close vs recvPkt
 	nl1, nl2 = _nodeLinkPipe(0, linkNoRecvSend)
 	c = xnewconn(nl1)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(c)
 	})
@@ -309,8 +327,8 @@ func TestNodeLink(t *testing.T) {
 	// Close vs sendPkt
 	nl1, nl2 = _nodeLinkPipe(0, linkNoRecvSend)
 	c = xnewconn(nl1)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		tdelay()
 		xclose(c)
 	})
@@ -324,14 +342,14 @@ func TestNodeLink(t *testing.T) {
 	// NodeLink.Close vs Conn.sendPkt/recvPkt
 	c11 := xnewconn(nl1)
 	c12 := xnewconn(nl1)
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		pkt, err := c11.recvPkt()
 		if !(pkt == nil && xconnError(err) == ErrLinkClosed) {
 			exc.Raisef("Conn.recvPkt() after NodeLink close: pkt = %v  err = %v", pkt, err)
 		}
 	})
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		pkt := c12.mkpkt(0, []byte("data"))
 		err := c12.sendPkt(pkt)
 		if xconnError(err) != ErrLinkClosed {
@@ -350,9 +368,9 @@ func TestNodeLink(t *testing.T) {
 	c21 := xnewconn(nl2)
 	c22 := xnewconn(nl2)
 	c23 := xnewconn(nl2)
-	wg = &errgroup.Group{}
+	wg = xsync.NewWorkGroup(bg)
 	var errRecv error
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		pkt, err := c21.recvPkt()
 		want1 := io.EOF           // if recvPkt wakes up due to peer close
 		want2 := io.ErrClosedPipe // if recvPkt wakes up due to sendPkt wakes up first and closes nl1
@@ -363,7 +381,7 @@ func TestNodeLink(t *testing.T) {
 
 		errRecv = cerr
 	})
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		pkt := c22.mkpkt(0, []byte("data"))
 		err := c22.sendPkt(pkt)
 		want := io.ErrClosedPipe // always this in both due to peer close or recvPkt waking up and closing nl2
@@ -372,7 +390,7 @@ func TestNodeLink(t *testing.T) {
 		}
 
 	})
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		conn, err := nl2.Accept()
 		if !(conn == nil && xlinkError(err) == ErrLinkDown) {
 			exc.Raisef("Accept after peer NodeLink shutdown: conn = %v  err = %v", conn, err)
@@ -466,9 +484,9 @@ func TestNodeLink(t *testing.T) {
 	// Conn accept + exchange
 	nl1, nl2 = nodeLinkPipe()
 	nl1.CloseAccept()
-	wg = &errgroup.Group{}
+	wg = xsync.NewWorkGroup(bg)
 	closed := make(chan int)
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		c := xaccept(nl2)
 
 		pkt := xrecvPkt(c)
@@ -560,7 +578,7 @@ func TestNodeLink(t *testing.T) {
 
 	// test 2 channels with replies coming in reversed time order
 	nl1, nl2 = nodeLinkPipe()
-	wg = &errgroup.Group{}
+	wg = xsync.NewWorkGroup(bg)
 	replyOrder := map[uint16]struct { // "order" in which to process requests
 		start chan struct{} // processing starts when start chan is ready
 		next  uint16        // after processing this switch to next
@@ -570,11 +588,11 @@ func TestNodeLink(t *testing.T) {
 	}
 	close(replyOrder[2].start)
 
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		for _ = range replyOrder {
 			c := xaccept(nl2)
 
-			gox(wg, func() {
+			gox(wg, func(_ context.Context) {
 				pkt := xrecvPkt(c)
 				n := packed.Ntoh16(pkt.Header().MsgCode)
 				x := replyOrder[n]
@@ -645,11 +663,13 @@ func xverifyMsg(msg1, msg2 proto.Msg) {
 }
 
 func TestRecv1Mode(t *testing.T) {
+	bg := context.Background()
+
 	// Send1
 	nl1, nl2 := nodeLinkPipe()
-	wg := &errgroup.Group{}
+	wg := xsync.NewWorkGroup(bg)
 	sync := make(chan int)
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		defer func() {
 			if e := recover(); e != nil {
 				panic(e)
@@ -685,8 +705,8 @@ func TestRecv1Mode(t *testing.T) {
 	xwait(wg)
 
 	// Recv1: further packets with same connid are rejected with "connection closed"
-	wg = &errgroup.Group{}
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(_ context.Context) {
 		c := xnewconn(nl2)
 
 		xSend(c, &proto.Ping{})
@@ -711,11 +731,11 @@ func TestRecv1Mode(t *testing.T) {
 // bug triggers under -race.
 func TestLightCloseVsLinkShutdown(t *testing.T) {
 	nl1, nl2 := nodeLinkPipe()
-	wg := &errgroup.Group{}
+	wg := xsync.NewWorkGroup(context.Background())
 
 	c := xnewconn(nl1)
 	nl1.shutdown()
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		c.lightClose()
 	})
 
@@ -1062,14 +1082,14 @@ func benchmarkLinkRTT(b *testing.B, l1, l2 *NodeLink) {
 func xlinkPipe(c1, c2 net.Conn) (*NodeLink, *NodeLink) {
 	var l1, l2 *NodeLink
 
-	wg := &errgroup.Group{}
-	gox(wg, func() {
-		l, err := _Handshake(context.Background(), c1, _LinkClient)
+	wg := xsync.NewWorkGroup(context.Background())
+	gox(wg, func(ctx context.Context) {
+		l, err := _Handshake(ctx, c1, _LinkClient)
 		exc.Raiseif(err)
 		l1 = l
 	})
-	gox(wg, func() {
-		l, err := _Handshake(context.Background(), c2, _LinkServer)
+	gox(wg, func(ctx context.Context) {
+		l, err := _Handshake(ctx, c2, _LinkServer)
 		exc.Raiseif(err)
 		l2 = l
 	})

go/neo/neonet/newlink_test.go

-// Copyright (C) 2016-2018  Nexedi SA and Contributors.
+// Copyright (C) 2016-2020  Nexedi SA and Contributors.
 //                          Kirill Smelkov <kirr@nexedi.com>
 //
 // This program is free software: you can Use, Study, Modify and Redistribute
@@ -25,9 +25,8 @@ import (
 	"net"
 	"testing"
 
-	"golang.org/x/sync/errgroup"
-
 	"lab.nexedi.com/kirr/go123/exc"
+	"lab.nexedi.com/kirr/go123/xsync"
 )
 
 func xhandshake(ctx context.Context, c net.Conn, version uint32) {
@@ -39,12 +38,12 @@ func TestHandshake(t *testing.T) {
 	bg := context.Background()
 	// handshake ok
 	p1, p2 := net.Pipe()
-	wg := &errgroup.Group{}
-	gox(wg, func() {
-		xhandshake(bg, p1, 1)
+	wg := xsync.NewWorkGroup(bg)
+	gox(wg, func(ctx context.Context) {
+		xhandshake(ctx, p1, 1)
 	})
-	gox(wg, func() {
-		xhandshake(bg, p2, 1)
+	gox(wg, func(ctx context.Context) {
+		xhandshake(ctx, p2, 1)
 	})
 	xwait(wg)
 	xclose(p1)
@@ -53,12 +52,12 @@ func TestHandshake(t *testing.T) {
 	// version mismatch
 	p1, p2 = net.Pipe()
 	var err1, err2 error
-	wg = &errgroup.Group{}
-	gox(wg, func() {
-		err1 = handshake(bg, p1, 1)
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(ctx context.Context) {
+		err1 = handshake(ctx, p1, 1)
 	})
-	gox(wg, func() {
-		err2 = handshake(bg, p2, 2)
+	gox(wg, func(ctx context.Context) {
+		err2 = handshake(ctx, p2, 2)
 	})
 	xwait(wg)
 	xclose(p1)
@@ -77,11 +76,11 @@ func TestHandshake(t *testing.T) {
 	// tx & rx problem
 	p1, p2 = net.Pipe()
 	err1, err2 = nil, nil
-	wg = &errgroup.Group{}
-	gox(wg, func() {
-		err1 = handshake(bg, p1, 1)
+	wg = xsync.NewWorkGroup(bg)
+	gox(wg, func(ctx context.Context) {
+		err1 = handshake(ctx, p1, 1)
 	})
-	gox(wg, func() {
+	gox(wg, func(_ context.Context) {
 		xclose(p2)
 	})
 	xwait(wg)
@@ -96,7 +95,8 @@ func TestHandshake(t *testing.T) {
 	// ctx cancel
 	p1, p2 = net.Pipe()
 	ctx, cancel := context.WithCancel(bg)
-	gox(wg, func() {
+	wg = xsync.NewWorkGroup(ctx)
+	gox(wg, func(ctx context.Context) {
 		err1 = handshake(ctx, p1, 1)
 	})
 	tdelay()