fix: roll back directChannel in orbitdb instance

D4ryl00 · D4ryl00 · commit 3301a16411f5 · 2022-12-06T23:50:46.000+01:00
Signed-off-by: D4ryl00 &lt;d4ryl00@gmail.com&gt;
diff --git a/baseorbitdb/events.go b/baseorbitdb/events.go
@@ -0,0 +1,18 @@
+package baseorbitdb
+
+import (
+	"berty.tech/go-orbit-db/iface"
+	"github.com/libp2p/go-libp2p/core/peer"
+)
+
+type EventExchangeHeads struct {
+	Peer    peer.ID
+	Message *iface.MessageExchangeHeads
+}
+
+func NewEventExchangeHeads(p peer.ID, msg *iface.MessageExchangeHeads) EventExchangeHeads {
+	return EventExchangeHeads{
+		Peer:    p,
+		Message: msg,
+	}
+}
diff --git a/baseorbitdb/events_handler.go b/baseorbitdb/events_handler.go
@@ -0,0 +1,26 @@
+package baseorbitdb
+
+import (
+	"context"
+	"fmt"
+
+	ipfslog "berty.tech/go-ipfs-log"
+	"berty.tech/go-orbit-db/iface"
+)
+
+func (o *orbitDB) handleEventExchangeHeads(ctx context.Context, e *iface.MessageExchangeHeads, store iface.Store) error {
+	untypedHeads := make([]ipfslog.Entry, len(e.Heads))
+	for i, h := range e.Heads {
+		untypedHeads[i] = h
+	}
+
+	o.logger.Debug(fmt.Sprintf("%s: Received %d heads for '%s':", o.PeerID().String(), len(untypedHeads), e.Address))
+
+	if len(untypedHeads) > 0 {
+		if err := store.Sync(ctx, untypedHeads); err != nil {
+			return fmt.Errorf("unable to sync heads: %w", err)
+		}
+	}
+
+	return nil
+}
diff --git a/baseorbitdb/orbitdb.go b/baseorbitdb/orbitdb.go
@@ -18,7 +18,9 @@ import (
 	"berty.tech/go-orbit-db/cache/cacheleveldown"
 	"berty.tech/go-orbit-db/iface"
 	_ "berty.tech/go-orbit-db/internal/buildconstraints" // fail for bad go version
-	"berty.tech/go-orbit-db/stores"
+	"berty.tech/go-orbit-db/messagemarshaler"
+	"berty.tech/go-orbit-db/pubsub"
+	"berty.tech/go-orbit-db/pubsub/oneonone"
 	"berty.tech/go-orbit-db/utils"
 	cid "github.com/ipfs/go-cid"
 	datastore "github.com/ipfs/go-datastore"
@@ -108,7 +110,7 @@ type orbitDB struct {
 	cache                 cache.Interface
 	logger                *zap.Logger
 	tracer                trace.Tracer
-	directChannelFactory  iface.DirectChannelFactory
+	directChannel         iface.DirectChannel
 	messageMarshaler      iface.MessageMarshaler
 
 	// emitters
@@ -178,17 +180,36 @@ func (o *orbitDB) setStore(address string, store iface.Store) {
 	o.stores[address] = store
 }
 
-func (o *orbitDB) closeAllStores() {
+func (o *orbitDB) deleteStore(address string) {
 	o.muStores.Lock()
 	defer o.muStores.Unlock()
 
+	delete(o.stores, address)
+}
+
+func (o *orbitDB) getStore(address string) (iface.Store, bool) {
+	o.muStores.RLock()
+	defer o.muStores.RUnlock()
+
+	store, ok := o.stores[address]
+
+	return store, ok
+}
+
+func (o *orbitDB) closeAllStores() {
+	stores := []Store{}
+
+	o.muStores.Lock()
 	for _, store := range o.stores {
+		stores = append(stores, store)
+	}
+	o.muStores.Unlock()
+
+	for _, store := range stores {
 		if err := store.Close(); err != nil {
 			o.logger.Error("unable to close store", zap.Error(err))
 		}
 	}
-
-	o.stores = map[string]Store{}
 }
 
 func (o *orbitDB) closeCache() {
@@ -200,6 +221,12 @@ func (o *orbitDB) closeCache() {
 	}
 }
 
+func (o *orbitDB) closeDirectConnections() {
+	if err := o.directChannel.Close(); err != nil {
+		o.logger.Error("unable to close connection", zap.Error(err))
+	}
+}
+
 func (o *orbitDB) closeKeyStore() {
 	o.muKeyStore.Lock()
 	defer o.muKeyStore.Unlock()
@@ -314,6 +341,16 @@ func newOrbitDB(ctx context.Context, is coreapi.CoreAPI, identity *idp.Identity,
 		options.EventBus = eventbus.NewBus()
 	}
 
+	if options.DirectChannelFactory == nil {
+		options.DirectChannelFactory = oneonone.NewChannelFactory(is)
+	}
+	directConnections, err := makeDirectChannel(ctx, options.EventBus, options.DirectChannelFactory, &iface.DirectChannelOptions{
+		Logger: options.Logger,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("unable to create a direct connection with peer: %w", err)
+	}
+
 	k, err := is.Key().Self(ctx)
 	if err != nil {
 		return nil, err
@@ -324,6 +361,10 @@ func newOrbitDB(ctx context.Context, is coreapi.CoreAPI, identity *idp.Identity,
 		options.PeerID = id
 	}
 
+	if options.MessageMarshaler == nil {
+		options.MessageMarshaler = &messagemarshaler.JSONMarshaler{}
+	}
+
 	if options.Cache == nil {
 		options.Cache = cacheleveldown.New(&cache.Options{Logger: options.Logger})
 	}
@@ -344,7 +385,7 @@ func newOrbitDB(ctx context.Context, is coreapi.CoreAPI, identity *idp.Identity,
 		directory:             *options.Directory,
 		eventBus:              options.EventBus,
 		stores:                map[string]Store{},
-		directChannelFactory:  options.DirectChannelFactory,
+		directChannel:         directConnections,
 		closeKeystore:         options.CloseKeystore,
 		storeTypes:            map[string]iface.StoreConstructor{},
 		accessControllerTypes: map[string]iface.AccessControllerConstructor{},
@@ -354,10 +395,15 @@ func newOrbitDB(ctx context.Context, is coreapi.CoreAPI, identity *idp.Identity,
 	}
 
 	// set new heads as stateful, so newly subscriber can replay last event in case they missed it
-	odb.emitters.newHeads, err = options.EventBus.Emitter(new(stores.EventExchangeHeads), eventbus.Stateful)
+	odb.emitters.newHeads, err = options.EventBus.Emitter(new(EventExchangeHeads), eventbus.Stateful)
 	if err != nil {
 		return nil, fmt.Errorf("unable to create global emitter: %w", err)
 	}
+
+	if err := odb.monitorDirectChannel(ctx, options.EventBus); err != nil {
+		return nil, fmt.Errorf("unable to monitor direct channel: %w", err)
+	}
+
 	return odb, nil
 }
 
@@ -430,9 +476,14 @@ func NewOrbitDB(ctx context.Context, ipfs coreapi.CoreAPI, options *NewOrbitDBOp
 
 func (o *orbitDB) Close() error {
 	o.closeAllStores()
+	o.closeDirectConnections()
 	o.closeCache()
 	o.closeKeyStore()
 
+	if err := o.emitters.newHeads.Close(); err != nil {
+		o.logger.Warn("unable to close emitter", zap.Error(err))
+	}
+
 	o.cancel()
 	return nil
 }
@@ -718,23 +769,35 @@ func (o *orbitDB) createStore(ctx context.Context, storeType string, parsedDBAdd
 		}
 	}
 
+	if options.Logger == nil {
+		options.Logger = o.logger
+	}
+
+	if options.CloseFunc == nil {
+		options.CloseFunc = func() {}
+	}
+	closeFunc := func() {
+		options.CloseFunc()
+		o.deleteStore(parsedDBAddress.String())
+	}
+
 	store, err := storeFunc(o.IPFS(), identity, parsedDBAddress, &iface.NewStoreOptions{
-		EventBus:             options.EventBus,
-		AccessController:     accessController,
-		Cache:                options.Cache,
-		Replicate:            options.Replicate,
-		Directory:            *options.Directory,
-		SortFn:               options.SortFn,
-		CacheDestroy:         func() error { return o.cache.Destroy(o.directory, parsedDBAddress) },
-		Logger:               o.logger,
-		Tracer:               o.tracer,
-		IO:                   options.IO,
-		StoreSpecificOpts:    options.StoreSpecificOpts,
-		PubSub:               o.pubsub,
-		MessageMarshaler:     o.messageMarshaler,
-		PeerID:               o.id,
-		DirectChannelFactory: o.directChannelFactory,
-		NewHeadsEmitter:      o.emitters.newHeads,
+		EventBus:          options.EventBus,
+		AccessController:  accessController,
+		Cache:             options.Cache,
+		Replicate:         options.Replicate,
+		Directory:         *options.Directory,
+		SortFn:            options.SortFn,
+		CacheDestroy:      func() error { return o.cache.Destroy(o.directory, parsedDBAddress) },
+		Logger:            options.Logger,
+		Tracer:            o.tracer,
+		IO:                options.IO,
+		StoreSpecificOpts: options.StoreSpecificOpts,
+		PubSub:            o.pubsub,
+		MessageMarshaler:  o.messageMarshaler,
+		PeerID:            o.id,
+		DirectChannel:     o.directChannel,
+		CloseFunc:         closeFunc,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("unable to instantiate store: %w", err)
@@ -749,4 +812,56 @@ func (o *orbitDB) EventBus() event.Bus {
 	return o.eventBus
 }
 
+func (o *orbitDB) monitorDirectChannel(ctx context.Context, bus event.Bus) error {
+	sub, err := bus.Subscribe(new(iface.EventPubSubPayload), eventbus.BufSize(128))
+	if err != nil {
+		return fmt.Errorf("unable to init pubsub subscriber: %w", err)
+	}
+
+	go func() {
+		for {
+			var e interface{}
+			select {
+			case <-ctx.Done():
+				return
+			case e = <-sub.Out():
+			}
+
+			evt := e.(iface.EventPubSubPayload)
+
+			msg := iface.MessageExchangeHeads{}
+			if err := o.messageMarshaler.Unmarshal(evt.Payload, &msg); err != nil {
+				o.logger.Error("unable to unmarshal message payload", zap.Error(err))
+				continue
+			}
+
+			store, ok := o.getStore(msg.Address)
+			if !ok {
+				o.logger.Error("unable to get store from address", zap.Error(err))
+				continue
+			}
+
+			if err := o.handleEventExchangeHeads(ctx, &msg, store); err != nil {
+				o.logger.Error("unable to handle pubsub payload", zap.Error(err))
+				continue
+			}
+
+			if err := o.emitters.newHeads.Emit(NewEventExchangeHeads(evt.Peer, &msg)); err != nil {
+				o.logger.Warn("unable to emit new heads", zap.Error(err))
+			}
+		}
+	}()
+
+	return nil
+}
+
+func makeDirectChannel(ctx context.Context, bus event.Bus, df iface.DirectChannelFactory, opts *iface.DirectChannelOptions) (iface.DirectChannel, error) {
+	emitter, err := pubsub.NewPayloadEmitter(bus)
+	if err != nil {
+		return nil, fmt.Errorf("unable to init pubsub emitter: %w", err)
+	}
+
+	return df(ctx, emitter, opts)
+}
+
 var _ BaseOrbitDB = &orbitDB{}
diff --git a/iface/interface.go b/iface/interface.go
@@ -53,6 +53,8 @@ type CreateDBOptions struct {
 	IO                      ipfslog.IO
 	Timeout                 time.Duration
 	MessageMarshaler        MessageMarshaler
+	Logger                  *zap.Logger
+	CloseFunc               func()
 	StoreSpecificOpts       interface{}
 }
 
@@ -345,12 +347,12 @@ type NewStoreOptions struct {
 	Logger                 *zap.Logger
 	Tracer                 trace.Tracer
 	IO                     ipfslog.IO
-	StoreSpecificOpts      interface{}
 	PubSub                 PubSubInterface
 	MessageMarshaler       MessageMarshaler
 	PeerID                 peer.ID
-	DirectChannelFactory   DirectChannelFactory
-	NewHeadsEmitter        event.Emitter
+	DirectChannel          DirectChannel
+	CloseFunc              func()
+	StoreSpecificOpts      interface{}
 }
 
 type DirectChannelOptions struct {
diff --git a/pubsub/directchannel/channel.go b/pubsub/directchannel/channel.go
@@ -82,6 +82,7 @@ func (d *directChannel) Connect(ctx context.Context, pid peer.ID) (err error) {
 // @NOTE(gfanton): we dont need this on direct channel
 // Close Closes the connection
 func (d *directChannel) Close() error {
+	d.host.RemoveStreamHandler(PROTOCOL)
 	return d.emitter.Close()
 }
 
diff --git a/stores/basestore/base_store.go b/stores/basestore/base_store.go
diff --git a/stores/events.go b/stores/events.go

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,7 @@ func (d *directChannel) Connect(ctx context.Context, pid peer.ID) (err error) {`
`82`	`82`	`// @NOTE(gfanton): we dont need this on direct channel`
`83`	`83`	`// Close Closes the connection`
`84`	`84`	`func (d *directChannel) Close() error {`
	`85`	`+ d.host.RemoveStreamHandler(PROTOCOL)`
`85`	`86`	`return d.emitter.Close()`
`86`	`87`	`}`
`87`	`88`