Add Persistence Layer on top of PubSub

fetch.go

@@ -0,0 +1,140 @@
+package namesys
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	ggio "github.com/gogo/protobuf/io"
+	"github.com/gogo/protobuf/proto"
+
+	"github.com/libp2p/go-libp2p-core/helpers"
+	"github.com/libp2p/go-libp2p-core/host"
+	"github.com/libp2p/go-libp2p-core/network"
+	"github.com/libp2p/go-libp2p-core/peer"
+	"github.com/libp2p/go-libp2p-core/protocol"
+
+	pb "github.com/libp2p/go-libp2p-pubsub-router/pb"
+)
+
+const FetchProtoID = protocol.ID("/libp2p/fetch/0.0.1")
+
+type fetchProtocol struct {
+	ctx  context.Context
+	host host.Host
+}
+
+type getValue func(key string) ([]byte, error)
+
+func newFetchProtocol(ctx context.Context, host host.Host, getData getValue) *fetchProtocol {
+	p := &fetchProtocol{ctx, host}
+
+	host.SetStreamHandler(FetchProtoID, func(s network.Stream) {
+		p.receive(s, getData)
+	})
+
+	return p
+}
+
+func (p *fetchProtocol) receive(s network.Stream, getData getValue) {
+	defer helpers.FullClose(s)
+
+	msg := &pb.FetchRequest{}
+	if err := readMsg(p.ctx, s, msg); err != nil {
+		log.Infof("error reading request from %s: %s", s.Conn().RemotePeer(), err)
+		s.Reset()
+		return
+	}
+
+	response, err := getData(msg.Identifier)
+	var respProto pb.FetchResponse
+
+	if err != nil {
+		respProto = pb.FetchResponse{Status: pb.FetchResponse_NOT_FOUND}
+	} else {
+		respProto = pb.FetchResponse{Data: response}
+	}
+
+	if err := writeMsg(p.ctx, s, &respProto); err != nil {
+		return
+	}
+}
+
+func (p *fetchProtocol) Fetch(ctx context.Context, pid peer.ID, key string) ([]byte, error) {
+	peerCtx, cancel := context.WithTimeout(ctx, time.Second*10)
+	defer cancel()
+
+	s, err := p.host.NewStream(peerCtx, pid, FetchProtoID)
+	if err != nil {
+		return nil, err
+	}
+	defer helpers.FullClose(s)
+
+	msg := &pb.FetchRequest{Identifier: key}
+
+	if err := writeMsg(ctx, s, msg); err != nil {
+		return nil, err
+	}
+	s.Close()
+
+	response := &pb.FetchResponse{}
+	if err := readMsg(ctx, s, response); err != nil {
+		return nil, err
+	}
+
+	switch response.Status {
+	case pb.FetchResponse_OK:
+		return response.Data, nil
+	case pb.FetchResponse_NOT_FOUND:
+		return nil, nil
+	default:
+		return nil, errors.New("fetch: received unknown status code")
+	}
+}
+
+func writeMsg(ctx context.Context, s network.Stream, msg proto.Message) error {
+	done := make(chan error, 1)
+	go func() {
+		wc := ggio.NewDelimitedWriter(s)
+
+		if err := wc.WriteMsg(msg); err != nil {
+			done <- err
+			return
+		}
+
+		done <- nil
+	}()
+
+	var retErr error
+	select {
+	case retErr = <-done:
+	case <-ctx.Done():
+		retErr = ctx.Err()
+	}
+
+	if retErr != nil {
+		s.Reset()
+		log.Infof("error writing response to %s: %s", s.Conn().RemotePeer(), retErr)
+	}
+	return retErr
+}
+
+func readMsg(ctx context.Context, s network.Stream, msg proto.Message) error {
+	done := make(chan error, 1)
+	go func() {
+		r := ggio.NewDelimitedReader(s, 1<<20)
+		if err := r.ReadMsg(msg); err != nil {
+			done <- err
+			return
+		}
+		done <- nil
+	}()
+
+	select {
+	case err := <-done:
+		return err
+	case <-ctx.Done():
+		s.Reset()
+		return ctx.Err()
+	}
+}

fetch_test.go

@@ -0,0 +1,109 @@
+package namesys
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/libp2p/go-libp2p-core/host"
+)
+
+func connect(t *testing.T, a, b host.Host) {
+	pinfo := a.Peerstore().PeerInfo(a.ID())
+	err := b.Connect(context.Background(), pinfo)
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+type datastore struct {
+	data map[string][]byte
+}
+
+func (d *datastore) Lookup(key string) ([]byte, error) {
+	v, ok := d.data[key]
+	if !ok {
+		return nil, errors.New("key not found")
+	}
+	return v, nil
+}
+
+func TestFetchProtocolTrip(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	hosts := newNetHosts(ctx, t, 2)
+	connect(t, hosts[0], hosts[1])
+
+	// wait for hosts to get connected
+	time.Sleep(time.Millisecond * 100)
+
+	d1 := &datastore{map[string][]byte{"key": []byte("value1")}}
+	h1 := newFetchProtocol(ctx, hosts[0], d1.Lookup)
+
+	d2 := &datastore{map[string][]byte{"key": []byte("value2")}}
+	h2 := newFetchProtocol(ctx, hosts[1], d2.Lookup)
+
+	fetchCheck(t, ctx, h1, h2, "key", []byte("value2"))
+	fetchCheck(t, ctx, h2, h1, "key", []byte("value1"))
+}
+
+func TestFetchProtocolNotFound(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	hosts := newNetHosts(ctx, t, 2)
+	connect(t, hosts[0], hosts[1])
+
+	// wait for hosts to get connected
+	time.Sleep(time.Millisecond * 100)
+
+	d1 := &datastore{map[string][]byte{"key": []byte("value1")}}
+	h1 := newFetchProtocol(ctx, hosts[0], d1.Lookup)
+
+	d2 := &datastore{make(map[string][]byte)}
+	h2 := newFetchProtocol(ctx, hosts[1], d2.Lookup)
+
+	fetchCheck(t, ctx, h1, h2, "key", nil)
+	fetchCheck(t, ctx, h2, h1, "key", []byte("value1"))
+}
+
+func TestFetchProtocolRepeated(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	hosts := newNetHosts(ctx, t, 2)
+	connect(t, hosts[0], hosts[1])
+
+	// wait for hosts to get connected
+	time.Sleep(time.Millisecond * 100)
+
+	d1 := &datastore{map[string][]byte{"key": []byte("value1")}}
+	h1 := newFetchProtocol(ctx, hosts[0], d1.Lookup)
+
+	d2 := &datastore{make(map[string][]byte)}
+	h2 := newFetchProtocol(ctx, hosts[1], d2.Lookup)
+
+	for i := 0; i < 10; i++ {
+		fetchCheck(t, ctx, h1, h2, "key", nil)
+		fetchCheck(t, ctx, h2, h1, "key", []byte("value1"))
+	}
+}
+
+func fetchCheck(t *testing.T, ctx context.Context,
+	requester *fetchProtocol, responder *fetchProtocol, key string, expected []byte) {
+	data, err := requester.Fetch(ctx, responder.host.ID(), key)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !bytes.Equal(data, expected) {
+		t.Fatalf("expected: %v, received: %v", string(expected), string(data))
+	}
+
+	if (data == nil && expected != nil) || (data != nil && expected == nil) {
+		t.Fatalf("expected []byte{} or nil and received the opposite")
+	}
+}

go.mod

@@ -1,14 +1,15 @@
 module github.com/libp2p/go-libp2p-pubsub-router
 
 require (
+	github.com/gogo/protobuf v1.2.1
 	github.com/ipfs/go-cid v0.0.2
 	github.com/ipfs/go-datastore v0.0.5
 	github.com/ipfs/go-ipfs-ds-help v0.0.1
 	github.com/ipfs/go-ipfs-util v0.0.1
 	github.com/ipfs/go-log v0.0.1
 	github.com/libp2p/go-libp2p-blankhost v0.1.1
 	github.com/libp2p/go-libp2p-core v0.0.1
-	github.com/libp2p/go-libp2p-pubsub v0.1.0
+	github.com/libp2p/go-libp2p-pubsub v0.1.1
 	github.com/libp2p/go-libp2p-record v0.1.0
 	github.com/libp2p/go-libp2p-routing-helpers v0.1.0
 	github.com/libp2p/go-libp2p-swarm v0.1.0

go.sum

@@ -95,8 +95,8 @@ github.com/libp2p/go-libp2p-peer v0.2.0 h1:EQ8kMjaCUwt/Y5uLgjT8iY2qg0mGUT0N1zUje
 github.com/libp2p/go-libp2p-peer v0.2.0/go.mod h1:RCffaCvUyW2CJmG2gAWVqwePwW7JMgxjsHm7+J5kjWY=
 github.com/libp2p/go-libp2p-peerstore v0.1.0 h1:MKh7pRNPHSh1fLPj8u/M/s/napdmeNpoi9BRy9lPN0E=
 github.com/libp2p/go-libp2p-peerstore v0.1.0/go.mod h1:2CeHkQsr8svp4fZ+Oi9ykN1HBb6u0MOvdJ7YIsmcwtY=
-github.com/libp2p/go-libp2p-pubsub v0.1.0 h1:SmQeMa7IUv5vadh0fYgYsafWCBA1sCy5d/68kIYqGcU=
-github.com/libp2p/go-libp2p-pubsub v0.1.0/go.mod h1:ZwlKzRSe1eGvSIdU5bD7+8RZN/Uzw0t1Bp9R1znpR/Q=
+github.com/libp2p/go-libp2p-pubsub v0.1.1 h1:phDnQvO3H3hAgaEEQi6yt3LILqIYVXaw05bxzezrEwQ=
+github.com/libp2p/go-libp2p-pubsub v0.1.1/go.mod h1:ZwlKzRSe1eGvSIdU5bD7+8RZN/Uzw0t1Bp9R1znpR/Q=
 github.com/libp2p/go-libp2p-record v0.1.0 h1:wHwBGbFzymoIl69BpgwIu0O6ta3TXGcMPvHUAcodzRc=
 github.com/libp2p/go-libp2p-record v0.1.0/go.mod h1:ujNc8iuE5dlKWVy6wuL6dd58t0n7xI4hAIl8pE6wu5Q=
 github.com/libp2p/go-libp2p-routing-helpers v0.1.0 h1:BaFvpyv8TyhCN7TihawTiKuzeu8/Pyw7ZnMA4IvqIN8=

pb/Makefile

@@ -0,0 +1,11 @@
+PB = $(wildcard *.proto)
+GO = $(PB:.proto=.pb.go)
+
+all: $(GO)
+
+%.pb.go: %.proto
+		protoc --proto_path="$(GOPATH)/src" --proto_path="." --gogofast_out=. $<
+
+clean:
+		rm -f *.pb.go
+		rm -f *.go