updated realtime broker to handle distributed deployment

Author: dstpierre

cache/cache.go

@@ -42,6 +42,14 @@ func (c *Cache) Set(key string, value string) error {
 	return nil
 }
 
+func (c *Cache) Inc(key string, by int64) (int64, error) {
+	return c.Rdb.IncrBy(c.Ctx, key, by).Result()
+}
+
+func (c *Cache) Dec(key string, by int64) (int64, error) {
+	return c.Rdb.DecrBy(c.Ctx, key, by).Result()
+}
+
 func (c *Cache) Subscribe(send chan internal.Command, token, channel string, close chan bool) {
 	pubsub := c.Rdb.Subscribe(c.Ctx, channel)
 

internal/data.go

@@ -39,6 +39,7 @@ const (
 	MsgTypeToken     = "token"
 	MsgTypeJoin      = "join"
 	MsgTypeJoined    = "joined"
+	MsgTypePresence  = "presence"
 	MsgTypeChanIn    = "chan_in"
 	MsgTypeChanOut   = "chan_out"
 	MsgTypeDBCreated = "db_created"

internal/pubsuber.go

@@ -0,0 +1,12 @@
+package internal
+
+// PubSuber contains functions to make realtime communication distributed
+type PubSuber interface {
+	Get(key string) (string, error)
+	Set(key string, value string) error
+	Inc(key string, by int64) (int64, error)
+	Dec(key string, by int64) (int64, error)
+	Subscribe(send chan Command, token, channel string, close chan bool)
+	Publish(msg Command) error
+	PublishDocument(channel, typ string, v interface{})
+}

main_test.go

@@ -95,7 +95,7 @@ func deleteAndSetupTestAccount() {
 		SBID:      acctID,
 		Name:      dbName,
 		Whitelist: []string{"localhost"},
-		Valid:     true,
+		IsActive:  true,
 	}
 
 	if _, err := sysDB.Collection("bases").InsertOne(ctx, base); err != nil {

realtime/broker.go

@@ -26,20 +26,23 @@ type Broker struct {
 	clients            map[chan internal.Command]string
 	ids                map[string]chan internal.Command
 	conf               map[string]context.Context
-	subscriptions      map[string][]string
+	subscriptions      map[string][]chan bool
 	validateAuth       Validator
+
+	pubsub internal.PubSuber
 }
 
-func NewBroker(v Validator) *Broker {
+func NewBroker(v Validator, pubsub internal.PubSuber) *Broker {
 	b := &Broker{
 		Broadcast:          make(chan internal.Command, 1),
 		newConnections:     make(chan ConnectionData),
 		closingConnections: make(chan chan internal.Command),
 		clients:            make(map[chan internal.Command]string),
 		ids:                make(map[string]chan internal.Command),
 		conf:               make(map[string]context.Context),
-		subscriptions:      make(map[string][]string),
+		subscriptions:      make(map[string][]chan bool),
 		validateAuth:       v,
+		pubsub:             pubsub,
 	}
 
 	go b.start()
@@ -85,6 +88,13 @@ func (b *Broker) unsub(c chan internal.Command) {
 		fmt.Println("cannot find connection id")
 	}
 
+	subs, ok := b.subscriptions[id]
+	if ok {
+		for _, ch := range subs {
+			ch <- true
+		}
+	}
+
 	delete(b.ids, id)
 }
 
@@ -140,12 +150,15 @@ func (b *Broker) Accept(w http.ResponseWriter, r *http.Request) {
 }
 
 func (b *Broker) getTargets(msg internal.Command) (sockets []chan internal.Command, payload internal.Command) {
+	var sender chan internal.Command
+
 	if msg.SID != internal.SystemID {
-		sender, ok := b.ids[msg.SID]
+		s, ok := b.ids[msg.SID]
 		if !ok {
 			fmt.Println("cannot find sender socket", msg.SID)
 			return
 		}
+		sender = s
 		sockets = append(sockets, sender)
 	}
 
@@ -167,15 +180,27 @@ func (b *Broker) getTargets(msg internal.Command) (sockets []chan internal.Comma
 
 		payload = internal.Command{Type: internal.MsgTypeToken, Data: msg.Data}
 	case internal.MsgTypeJoin:
-		members, ok := b.subscriptions[msg.Data]
+		subs, ok := b.subscriptions[msg.SID]
 		if !ok {
-			members = make([]string, 0)
+			subs = make([]chan bool, 0)
 		}
 
-		members = append(members, msg.SID)
-		b.subscriptions[msg.Data] = members
+		closesub := make(chan bool)
+
+		subs = append(subs, closesub)
+		b.subscriptions[msg.SID] = subs
+
+		go b.pubsub.Subscribe(sender, msg.Token, msg.Data, closesub)
 
 		payload = internal.Command{Type: internal.MsgTypeJoined, Data: msg.Data}
+	case internal.MsgTypePresence:
+		v, err := b.pubsub.Get(msg.Data)
+		if err != nil {
+			//TODO: Make sure it's because the channel key does not exists
+			v = "0"
+		}
+
+		payload = internal.Command{Type: internal.MsgTypePresence, Data: v}
 	case internal.MsgTypeChanIn:
 		if len(msg.Channel) == 0 {
 			payload = internal.Command{Type: internal.MsgTypeError, Data: "no channel was specified"}
@@ -188,7 +213,8 @@ func (b *Broker) getTargets(msg internal.Command) (sockets []chan internal.Comma
 			return
 		}
 
-		go b.Publish(msg, msg.Channel)
+		go b.pubsub.Publish(msg)
+		//go b.Publish(msg, msg.Channel)
 
 		payload = internal.Command{Type: internal.MsgTypeOk}
 	default:
@@ -198,24 +224,3 @@ func (b *Broker) getTargets(msg internal.Command) (sockets []chan internal.Comma
 
 	return
 }
-
-// Publish sends a message to all socket in that channel
-func (b *Broker) Publish(msg internal.Command, channel string) {
-	if msg.Type == internal.MsgTypeChanIn {
-		msg.Type = internal.MsgTypeChanOut
-	}
-
-	members, ok := b.subscriptions[channel]
-	if !ok {
-		return
-	}
-
-	for _, sid := range members {
-		c, ok := b.ids[sid]
-		if !ok {
-			continue
-		}
-
-		c <- msg
-	}
-}

server.go

@@ -19,6 +19,7 @@ import (
 	"staticbackend/cache"
 
 	"github.com/stripe/stripe-go/v71"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.mongodb.org/mongo-driver/mongo/readpref"
@@ -53,11 +54,26 @@ func Start(dbHost, port string) {
 
 	// Server Send Event, alternative to websocket
 	b := realtime.NewBroker(func(ctx context.Context, key string) (string, error) {
+		//TODO: Experimental, let un-authenticated user connect
+		// useful for an Intercom-like SaaS I'm building.
+		if strings.HasPrefix(key, "__tmp__experimental_public_19378246_") {
+			// let's create the most minimal authentication possible
+			a := internal.Auth{
+				AccountID: primitive.NewObjectID(),
+				UserID:    primitive.NewObjectID(),
+				Email:     "exp@tmp.com",
+				Role:      0,
+			}
+
+			internal.Tokens[key] = a
+			return key, nil
+		}
+
 		if _, err := middleware.ValidateAuthKey(client, ctx, key); err != nil {
 			return "", err
 		}
 		return key, nil
-	})
+	}, volatile)
 
 	database := &Database{
 		client: client,