Fix persistent connection request handling (#10)

Author: dvmazur

persistent_stream.go

@@ -45,67 +45,63 @@ func (d *Daemon) handlePersistentConn(r ggio.Reader, unsafeW ggio.WriteCloser) {
 			return
 		}
 
-		callID, err := uuid.FromBytes(req.CallId)
-		if err != nil {
-			log.Debugw("bad call id: ", "error", err)
-			continue
+		go d.handlePersistentConnRequest(req, w, &streamHandlers)
+	}
+}
+
+func (d *Daemon) handlePersistentConnRequest(req pb.PersistentConnectionRequest, w ggio.WriteCloser, streamHandlers *[]string) {
+	callID, err := uuid.FromBytes(req.CallId)
+	if err != nil {
+		log.Debugw("bad call id: ", "error", err)
+		return
+	}
+
+	switch req.Message.(type) {
+	case *pb.PersistentConnectionRequest_AddUnaryHandler:
+		resp := d.doAddUnaryHandler(w, callID, req.GetAddUnaryHandler())
+
+		d.mx.Lock()
+		if _, ok := resp.Message.(*pb.PersistentConnectionResponse_DaemonError); !ok {
+			*streamHandlers = append(
+				*streamHandlers,
+				*req.GetAddUnaryHandler().Proto,
+			)
 		}
+		d.mx.Unlock()
 
-		switch req.Message.(type) {
-		case *pb.PersistentConnectionRequest_AddUnaryHandler:
-			go func() {
-				resp := d.doAddUnaryHandler(w, callID, req.GetAddUnaryHandler())
-
-				d.mx.Lock()
-				if _, ok := resp.Message.(*pb.PersistentConnectionResponse_DaemonError); !ok {
-					streamHandlers = append(
-						streamHandlers,
-						*req.GetAddUnaryHandler().Proto,
-					)
-				}
-				d.mx.Unlock()
-
-				if err := w.WriteMsg(resp); err != nil {
-					log.Debugw("error reading message", "error", err)
-					return
-				}
-			}()
-
-		case *pb.PersistentConnectionRequest_CallUnary:
-			go func() {
-				ctx, cancel := context.WithCancel(context.Background())
-				d.cancelUnary.Store(callID, cancel)
-				defer cancel()
-
-				defer d.cancelUnary.Delete(callID)
-
-				resp := d.doUnaryCall(ctx, callID, &req)
-
-				if err := w.WriteMsg(resp); err != nil {
-					log.Debugw("error reading message", "error", err)
-					return
-				}
-			}()
-
-		case *pb.PersistentConnectionRequest_UnaryResponse:
-			go func() {
-				resp := d.sendReponseToRemote(&req)
-				if err := w.WriteMsg(resp); err != nil {
-					log.Debugw("error reading message", "error", err)
-					return
-				}
-			}()
-
-		case *pb.PersistentConnectionRequest_Cancel:
-			go func() {
-				cf, found := d.cancelUnary.Load(callID)
-				if !found {
-					return
-				}
-
-				cf.(context.CancelFunc)()
-			}()
+		if err := w.WriteMsg(resp); err != nil {
+			log.Debugw("error reading message", "error", err)
+			return
 		}
+
+	case *pb.PersistentConnectionRequest_CallUnary:
+		ctx, cancel := context.WithCancel(context.Background())
+		d.cancelUnary.Store(callID, cancel)
+		defer cancel()
+
+		defer d.cancelUnary.Delete(callID)
+
+		resp := d.doUnaryCall(ctx, callID, &req)
+
+		if err := w.WriteMsg(resp); err != nil {
+			log.Debugw("error reading message", "error", err)
+			return
+		}
+
+	case *pb.PersistentConnectionRequest_UnaryResponse:
+		resp := d.sendReponseToRemote(&req)
+		if err := w.WriteMsg(resp); err != nil {
+			log.Debugw("error reading message", "error", err)
+			return
+		}
+
+	case *pb.PersistentConnectionRequest_Cancel:
+		cf, found := d.cancelUnary.Load(callID)
+		if !found {
+			return
+		}
+
+		cf.(context.CancelFunc)()
 	}
 }
 

test/unary_handler_test.go

@@ -6,13 +6,60 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/libp2p/go-libp2p-core/protocol"
 	"github.com/libp2p/go-libp2p-daemon/p2pclient"
 )
 
+func TestConcurrentCalls(t *testing.T) {
+	_, p1, cancel1 := createDaemonClientPair(t)
+	_, p2, cancel2 := createDaemonClientPair(t)
+
+	defer func() {
+		cancel1()
+		cancel2()
+	}()
+
+	peer1ID, peer1Addrs, err := p1.Identify()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := p2.Connect(peer1ID, peer1Addrs); err != nil {
+		t.Fatal(err)
+	}
+
+	var proto protocol.ID = "sqrt"
+	if err := p1.AddUnaryHandler(proto, sqrtHandler); err != nil {
+		t.Fatal(err)
+	}
+
+	count := 100
+
+	var wg sync.WaitGroup
+	var m sync.Map
+	wg.Add(count)
+
+	for i := 0; i < count; i++ {
+		go func(i int) {
+			defer wg.Done()
+
+			reply, err := p2.CallUnaryHandler(context.Background(), peer1ID, proto, float64Bytes(float64(i)))
+			if err != nil {
+				panic(err)
+			}
+
+			if _, loaded := m.LoadOrStore(float64FromBytes(reply), ""); loaded {
+				panic(err)
+			}
+		}(i)
+	}
+
+	wg.Wait()
+}
+
 func TestUnaryCalls(t *testing.T) {
 	_, p1, cancel1 := createDaemonClientPair(t)
 	_, p2, cancel2 := createDaemonClientPair(t)