Persistent connections and unary handlers (#3)

Author: dvmazur

.gitignore

@@ -1,2 +1,3 @@
 .DS_Store
 .idea/
+.vscode/

conn.go

@@ -146,6 +146,10 @@ func (d *Daemon) handleConn(c net.Conn) {
 				return
 			}
 
+		case pb.Request_PERSISTENT_CONN_UPGRADE:
+			d.handlePersistentConn(r, w)
+			return
+
 		default:
 			log.Debugw("unexpected request type", "type", req.GetType())
 			return

daemon.go

@@ -3,6 +3,7 @@ package p2pd
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"os"
 	"sync"
@@ -39,12 +40,28 @@ type Daemon struct {
 	handlers map[protocol.ID]ma.Multiaddr
 	// closed is set when the daemon is shutting down
 	closed bool
+
+	registeredUnaryProtocols map[protocol.ID]bool
+
+	// callID (int64) to chan *pb.PersistentConnectionResponse
+	// used to return responses to goroutines awating them
+	responseWaiters sync.Map
+	// callID (int64) to chan context.CancelFunc
+	// used to cancel request handlers
+	cancelUnary sync.Map
+
+	// this sync.Once ensures the goroutine awaiting deamon termination is
+	// only run once
+	terminateOnce        sync.Once
+	terminateWG          sync.WaitGroup
+	cancelTerminateTimer context.CancelFunc
 }
 
 func NewDaemon(ctx context.Context, maddr ma.Multiaddr, dhtMode string, opts ...libp2p.Option) (*Daemon, error) {
 	d := &Daemon{
-		ctx:      ctx,
-		handlers: make(map[protocol.ID]ma.Multiaddr),
+		ctx:                      ctx,
+		handlers:                 make(map[protocol.ID]ma.Multiaddr),
+		registeredUnaryProtocols: make(map[protocol.ID]bool),
 	}
 
 	if dhtMode != "" {
@@ -71,7 +88,6 @@ func NewDaemon(ctx context.Context, maddr ma.Multiaddr, dhtMode string, opts ...
 	}
 	d.listener = l
 
-	go d.listen()
 	go d.trapSignals()
 
 	return d, nil
@@ -134,10 +150,10 @@ func (d *Daemon) Addrs() []ma.Multiaddr {
 	return d.host.Addrs()
 }
 
-func (d *Daemon) listen() {
+func (d *Daemon) Serve() error {
 	for {
 		if d.isClosed() {
-			return
+			return nil
 		}
 
 		c, err := d.listener.Accept()
@@ -191,3 +207,22 @@ func (d *Daemon) Close() error {
 
 	return merr.ErrorOrNil()
 }
+
+func (d *Daemon) awaitTermination() {
+	d.terminateWG.Wait()
+	d.Close()
+}
+
+func (d *Daemon) KillOnTimeout(timeout time.Duration) {
+	go func() {
+		ctx, cancel := context.WithCancel(d.ctx)
+		d.cancelTerminateTimer = cancel
+
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.NewTimer(timeout).C:
+			d.Close()
+		}
+	}()
+}

go.mod

@@ -5,6 +5,7 @@ go 1.15
 require (
 	github.com/davidlazar/go-crypto v0.0.0-20200604182044-b73af7476f6c // indirect
 	github.com/gogo/protobuf v1.3.2
+	github.com/google/uuid v1.3.0
 	github.com/hashicorp/go-multierror v1.1.0
 	github.com/ipfs/go-cid v0.0.7
 	github.com/ipfs/go-log v1.0.5

go.sum

@@ -184,8 +184,9 @@ github.com/google/pprof v0.0.0-20181206194817-3ea8567a2e57/go.mod h1:zfwlbNMJ+OI
 github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
 github.com/google/uuid v1.0.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/gax-go v2.0.0+incompatible/go.mod h1:SFVmujtThgffbyetf+mdk2eWhX2bMyUtNHzFKcPA9HY=
 github.com/googleapis/gax-go/v2 v2.0.3/go.mod h1:LLvjysVCY1JZeum8Z6l8qUty8fiNwE08qbEPm1M08qg=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=

internal/utils/safe_writer.go

@@ -0,0 +1,29 @@
+package utils
+
+import (
+	"sync"
+
+	ggio "github.com/gogo/protobuf/io"
+	"github.com/gogo/protobuf/proto"
+)
+
+func NewSafeWriter(w ggio.WriteCloser) *safeWriter {
+	return &safeWriter{w: w}
+}
+
+type safeWriter struct {
+	w ggio.WriteCloser
+	m sync.Mutex
+}
+
+func (sw *safeWriter) WriteMsg(msg proto.Message) error {
+	sw.m.Lock()
+	defer sw.m.Unlock()
+	return sw.w.WriteMsg(msg)
+}
+
+func (sw *safeWriter) Close() error {
+	sw.m.Lock()
+	defer sw.m.Unlock()
+	return sw.w.Close()
+}

p2pclient/p2pclient.go

@@ -26,6 +26,13 @@ type Client struct {
 
 	mhandlers sync.Mutex
 	handlers  map[string]StreamHandlerFunc
+
+	openPersistentConn   sync.Once
+	persistentConnWriter ggio.WriteCloser
+
+	// callID (uuid.UUID) -> persistentConnectionFuture
+	callFutures   sync.Map
+	unaryHandlers sync.Map
 }
 
 // NewClient creates a new libp2p daemon client, connecting to a daemon

p2pclient/streams.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 
+	"github.com/hashicorp/go-multierror"
 	"github.com/libp2p/go-libp2p-core/peer"
 
 	ggio "github.com/gogo/protobuf/io"
@@ -125,11 +126,17 @@ func (c *Client) NewStream(peer peer.ID, protos []string) (*StreamInfo, io.ReadW
 
 // Close stops the listener address.
 func (c *Client) Close() error {
+	merr := &multierror.Error{}
+
 	if c.listener != nil {
-		err := c.listener.Close()
-		return err
+		multierror.Append(merr, c.listener.Close())
 	}
-	return nil
+
+	if c.persistentConnWriter != nil {
+		multierror.Append(merr, c.persistentConnWriter.Close())
+	}
+
+	return merr.ErrorOrNil()
 }
 
 func (c *Client) streamDispatcher() {