wehub-resource-sync
452 lines
15 KiB
Go
452 lines
15 KiB
Go
package apis
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"math/rand"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
"github.com/jackc/pgx/v5/stdlib"
|
|
"github.com/pocketbase/dbx"
|
|
"github.com/pocketbase/pocketbase/core"
|
|
"github.com/pocketbase/pocketbase/tools/security"
|
|
"github.com/pocketbase/pocketbase/tools/subscriptions"
|
|
)
|
|
|
|
type IRealtimeBridge interface {
|
|
// App returns the current app instance.
|
|
App() core.App
|
|
|
|
// SelfChannelId returns the channel id of the current instance.
|
|
SelfChannelId() string
|
|
|
|
// SendViaBridge sends a message to a remote client via the bridge.
|
|
SendViaBridge(channelId string, clientId string, message subscriptions.Message)
|
|
}
|
|
|
|
var _ IRealtimeBridge = (*RealtimeBridge)(nil)
|
|
|
|
// RealtimeBridge leverage PostgresSQL's LISTEN/NOTIFY feature to synchronize
|
|
// realtime information between different instances of the pocketbase server.
|
|
type RealtimeBridge struct {
|
|
channelId string
|
|
app core.App
|
|
}
|
|
|
|
var RealtimeBridgeInstanceKey = "realtime_bridge_instance"
|
|
|
|
func bindRealtimeBridge(app core.App) {
|
|
if !core.IsPostgres(app) {
|
|
return
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
bridge := &RealtimeBridge{
|
|
channelId: genChannelId(),
|
|
app: app,
|
|
}
|
|
|
|
app.OnServe().BindFunc(func(e *core.ServeEvent) error {
|
|
bridge.mustCreateTables()
|
|
go bridge.heartbeatLoop(ctx)
|
|
go bridge.listenSharedBridgeChannelLoop(ctx)
|
|
go bridge.listenBridgeMessagesLoop(ctx)
|
|
|
|
// allow bridge instance to be accessible from via app.Store()
|
|
app.Store().Set(RealtimeBridgeInstanceKey, bridge)
|
|
|
|
return e.Next()
|
|
})
|
|
|
|
app.OnTerminate().BindFunc(func(e *core.TerminateEvent) error {
|
|
cancel() // stopping all background goroutines
|
|
return e.Next()
|
|
})
|
|
|
|
// Special fixes for horizontally deployed pocketbase instances.
|
|
app.OnCollectionCreateRequest().BindFunc(bridge.broadcastCollectionChanged)
|
|
app.OnCollectionUpdateRequest().BindFunc(bridge.broadcastCollectionChanged)
|
|
app.OnCollectionDeleteRequest().BindFunc(bridge.broadcastCollectionChanged)
|
|
app.OnSettingsUpdateRequest().BindFunc(bridge.broadcastSettingsUpdated)
|
|
}
|
|
|
|
// listenSharedBridgeChannelLoop listens to the shared bridge channel.
|
|
// It is a common communication channel between all pocketbase instances.
|
|
// Currently, it has two purposes:
|
|
// 1. Listen upsert, delete events in _realtimeClients table.
|
|
// 2. Listen collection_updated and settings_updated events.
|
|
func (t *RealtimeBridge) listenSharedBridgeChannelLoop(ctx context.Context) {
|
|
loopOnNotification(ctx, t.app, "shared_bridge_channel", func() {
|
|
// When it connected to the stream, we need to reload all subscriptions
|
|
// to make sure that we have the latest state.
|
|
t.fullRefreshSubscriptions()
|
|
|
|
// Reload collections and settings in case someone else updated them and
|
|
// this instance somehow didn't get the changes.
|
|
_ = t.app.ReloadCachedCollections()
|
|
_ = t.app.ReloadSettings()
|
|
}, func(notification *pgconn.Notification) {
|
|
if t.app.IsDev() {
|
|
fmt.Println("PID:", notification.PID, "Channel:", notification.Channel, "Payload:", notification.Payload)
|
|
}
|
|
|
|
messageType, messagePayload, ok := split2(notification.Payload, "|")
|
|
if !ok {
|
|
fmt.Fprintln(os.Stderr, "Invalid notification payload:", notification.Payload)
|
|
return
|
|
}
|
|
|
|
switch messageType {
|
|
case "subscription_upsert":
|
|
subscriptionJson, authRecordJson, ok := split2(messagePayload, "|")
|
|
if !ok {
|
|
fmt.Fprintln(os.Stderr, "Invalid subscriptionChange payload:", messagePayload)
|
|
return
|
|
}
|
|
|
|
var subscription ClientSubscription
|
|
err := json.Unmarshal([]byte(subscriptionJson), &subscription)
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error unmarshalling notification payload:", err)
|
|
os.Exit(1)
|
|
}
|
|
if subscription.UpdatedByChannelId == t.SelfChannelId() {
|
|
// Since the notification is broadcast to all channels, we need to skip the
|
|
// ones that are sent by the current channel.
|
|
return
|
|
}
|
|
|
|
var client IBridgedClient
|
|
if c, err := t.app.SubscriptionsBroker().ClientById(subscription.ClientId); err == nil {
|
|
client = c.(IBridgedClient)
|
|
} else {
|
|
client = NewBridgedClient(t, &subscription)
|
|
t.app.SubscriptionsBroker().Register(client)
|
|
}
|
|
client.ReceiveChanges(&subscription, authRecordJson)
|
|
case "subscription_delete":
|
|
clientId := messagePayload
|
|
t.app.SubscriptionsBroker().Unregister(clientId)
|
|
case "subscription_channel_offline":
|
|
channelId := messagePayload
|
|
// unregister all remote clients in that channel
|
|
for _, c := range t.app.SubscriptionsBroker().Clients() {
|
|
if syncClient, ok := c.(IBridgedClient); ok && syncClient.ClientSubscription().ChannelId == channelId {
|
|
t.app.SubscriptionsBroker().Unregister(syncClient.Id())
|
|
}
|
|
}
|
|
case "collection_updated":
|
|
_ = t.app.ReloadCachedCollections()
|
|
case "settings_updated":
|
|
_ = t.app.ReloadSettings()
|
|
default:
|
|
fmt.Fprintln(os.Stderr, "Unknown change type:", messageType)
|
|
panic("Unknown change type: " + messageType)
|
|
}
|
|
})
|
|
}
|
|
|
|
func (t *RealtimeBridge) SendViaBridge(channelId string, clientId string, message subscriptions.Message) {
|
|
if channelId == t.SelfChannelId() {
|
|
fmt.Fprintln(os.Stderr, "Cannot send bridged message to self channel:", channelId)
|
|
panic("Cannot send bridged message to self channel")
|
|
}
|
|
_, err := t.app.DB().NewQuery(`
|
|
SELECT pg_notify({:channelId}, {:payload})
|
|
`).Bind(dbx.Params{
|
|
"channelId": channelId,
|
|
"payload": clientId + "|" + message.Name + "|" + string(message.Data),
|
|
}).Execute()
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error sending notification:", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
func (b *RealtimeBridge) listenBridgeMessagesLoop(ctx context.Context) {
|
|
loopOnNotification(ctx, b.app, b.SelfChannelId(), nil, func(notification *pgconn.Notification) {
|
|
if b.app.IsDev() {
|
|
fmt.Println("PID:", notification.PID, "Channel:", notification.Channel, "Payload:", notification.Payload)
|
|
}
|
|
|
|
clientId, messageName, messageData, ok := split3(notification.Payload, "|")
|
|
if !ok {
|
|
fmt.Fprintln(os.Stderr, "Invalid notification payload:", notification.Payload)
|
|
return
|
|
}
|
|
var client IBridgedClient
|
|
if c, err := b.app.SubscriptionsBroker().ClientById(clientId); err == nil {
|
|
client = c.(IBridgedClient)
|
|
} else {
|
|
fmt.Fprintln(os.Stderr, "Client not found, it may be already disconnected:", clientId)
|
|
return
|
|
}
|
|
// Message is send to the wrong channel.
|
|
// Eg: a message was supposed to be sent to local clientA which is in channelA.
|
|
// But somehow the message was sent to remote clientA in channelB.
|
|
if client.IsRemoteClient() {
|
|
fmt.Fprintln(os.Stderr, "Message is sent to the wrong channel:", notification.Payload)
|
|
return
|
|
}
|
|
client.Send(subscriptions.Message{
|
|
Name: messageName,
|
|
Data: []byte(messageData),
|
|
})
|
|
})
|
|
}
|
|
|
|
var pgTypes = pgtype.NewMap()
|
|
|
|
// reload all remote realtime subscriptions
|
|
func (t *RealtimeBridge) fullRefreshSubscriptions() {
|
|
rows, err := t.app.DB().NewQuery(`
|
|
SELECT "clientId", "channelId", "subscriptions", "authCollectionRef", "authRecordRef", "updatedByChannelId"
|
|
FROM "_realtimeClients"
|
|
WHERE "updatedByChannelId" != {:selfChannelId}
|
|
`).Bind(dbx.Params{
|
|
"selfChannelId": t.SelfChannelId(),
|
|
}).Rows()
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error loading subscriptions:", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
clientsToRemove := make(map[string]any)
|
|
for _, c := range t.app.SubscriptionsBroker().Clients() {
|
|
if client, ok := c.(IBridgedClient); ok && client.IsRemoteClient() {
|
|
clientsToRemove[client.Id()] = nil
|
|
}
|
|
}
|
|
|
|
for rows.Next() {
|
|
var subscription ClientSubscription
|
|
err := rows.Scan(
|
|
&subscription.ClientId,
|
|
&subscription.ChannelId,
|
|
pgTypes.SQLScanner(&subscription.Subscriptions),
|
|
&subscription.AuthCollectionRef,
|
|
&subscription.AuthRecordRef,
|
|
&subscription.UpdatedByChannelId,
|
|
)
|
|
if err != nil {
|
|
data := make(dbx.NullStringMap, 0)
|
|
err := rows.ScanMap(data)
|
|
fmt.Fprintln(os.Stderr, "Error scanning subscription:", err)
|
|
continue
|
|
}
|
|
var client IBridgedClient
|
|
if c, err := t.app.SubscriptionsBroker().ClientById(subscription.ClientId); err == nil {
|
|
client = c.(IBridgedClient)
|
|
} else {
|
|
client = NewBridgedClient(t, &subscription)
|
|
t.app.SubscriptionsBroker().Register(client)
|
|
}
|
|
client.ReceiveChanges(&subscription, "")
|
|
delete(clientsToRemove, client.Id())
|
|
}
|
|
for clientId := range clientsToRemove {
|
|
t.app.SubscriptionsBroker().Unregister(clientId)
|
|
}
|
|
}
|
|
|
|
func (t *RealtimeBridge) mustCreateTables() {
|
|
_, err := t.app.DB().NewQuery(`
|
|
CREATE TABLE IF NOT EXISTS "_realtimeChannels" (
|
|
"channelId" TEXT PRIMARY KEY,
|
|
"validUntil" TIMESTAMP NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS "_realtimeClients" (
|
|
"clientId" TEXT NOT NULL PRIMARY KEY,
|
|
"channelId" TEXT NOT NULL,
|
|
"subscriptions" TEXT[] NOT NULL,
|
|
"authCollectionRef" TEXT NOT NULL DEFAULT '',
|
|
"authRecordRef" TEXT NOT NULL DEFAULT '',
|
|
"updatedByChannelId" TEXT NOT NULL DEFAULT ''
|
|
);
|
|
CREATE INDEX IF NOT EXISTS "_realtimeClients_channelId_idx" ON "_realtimeClients" ("channelId");
|
|
`).Execute()
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error creating tables tables:", err)
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
// heartbeatLoop periodically updates its status in the _realtimeChannels table
|
|
// to tell other pocketbase instances that it is still alive.
|
|
// It also helps to broadcast the subscription_channel_offline event to all pocketbase instances
|
|
// if any of them is offline.
|
|
func (t *RealtimeBridge) heartbeatLoop(ctx context.Context) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
fmt.Fprintln(os.Stderr, "Stopping realtime sync heartbeat loop.")
|
|
return
|
|
default:
|
|
if t.app.DB() == nil {
|
|
fmt.Fprintln(os.Stderr, "App is not initialized or stopped, stopping realtime sync heartbeat loop.")
|
|
break
|
|
}
|
|
_, err := t.app.DB().NewQuery(`
|
|
WITH
|
|
insert_operation AS (
|
|
INSERT INTO "_realtimeChannels" ("channelId", "validUntil")
|
|
VALUES ({:channelId}, now() + interval '40 seconds')
|
|
ON CONFLICT ("channelId") DO UPDATE
|
|
SET "validUntil" = EXCLUDED."validUntil"
|
|
),
|
|
deleted_channels AS (
|
|
DELETE FROM "_realtimeChannels"
|
|
WHERE "validUntil" < now()
|
|
RETURNING "channelId"
|
|
),
|
|
_ AS (
|
|
DELETE FROM "_realtimeClients"
|
|
WHERE "channelId" IN (SELECT "channelId" FROM deleted_channels)
|
|
)
|
|
SELECT pg_notify('shared_bridge_channel', 'subscription_channel_offline|' || "channelId") FROM deleted_channels;
|
|
`).Bind(dbx.Params{
|
|
"channelId": t.channelId,
|
|
}).Execute()
|
|
if err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error sending heartbeat:", err)
|
|
time.Sleep(time.Second * 1)
|
|
continue
|
|
}
|
|
jitter := time.Second * time.Duration(rand.Intn(5))
|
|
time.Sleep(time.Second*30 - jitter)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (t *RealtimeBridge) SelfChannelId() string {
|
|
return t.channelId
|
|
}
|
|
|
|
func (t *RealtimeBridge) App() core.App {
|
|
return t.app
|
|
}
|
|
|
|
func genChannelId() string {
|
|
hostname, _ := os.Hostname()
|
|
randstr, _ := security.RandomStringByRegex(`[a-z0-9]{5}`)
|
|
channelId := "c" + "_" + hostname + "_" + randstr
|
|
|
|
// Normalize the channelId to be a valid Postgres identifier
|
|
// Only allow lowercase letters, numbers and underscores
|
|
channelId = strings.Map(func(r rune) rune {
|
|
if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' {
|
|
return r
|
|
}
|
|
if r >= 'A' && r <= 'Z' {
|
|
return r + ('a' - 'A') // Convert to lowercase
|
|
}
|
|
return '_'
|
|
}, channelId)
|
|
|
|
return channelId
|
|
}
|
|
|
|
func loopOnNotification(ctx context.Context, app core.App, channel string, afterListenFunc func(), handler func(*pgconn.Notification)) {
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
fmt.Fprintln(os.Stderr, "Stopping PostgreSQL stream listener loop on channel:", channel)
|
|
return
|
|
default:
|
|
err := runPgxCommand(ctx, app, func(pgxConn *pgx.Conn) error {
|
|
_, err := pgxConn.Exec(ctx, fmt.Sprintf("LISTEN %s", channel))
|
|
if err != nil {
|
|
return fmt.Errorf("error listening on channel %s: %w", channel, err)
|
|
}
|
|
|
|
if afterListenFunc != nil {
|
|
afterListenFunc()
|
|
}
|
|
|
|
for {
|
|
notification, err := pgxConn.WaitForNotification(ctx)
|
|
if err != nil {
|
|
if err == context.Canceled {
|
|
app.Logger().Info("Context was canceled, exiting the loop on channel", "channel", channel)
|
|
return nil // context was canceled, exit the loop
|
|
}
|
|
return fmt.Errorf("error waiting for notification on channel %s: %w", channel, err)
|
|
}
|
|
handler(notification)
|
|
}
|
|
})
|
|
if err != nil {
|
|
app.Logger().Error("Error in PostgreSQL stream listener loop on channel", "channel", channel, "error", err)
|
|
fmt.Fprintln(os.Stderr, "Error in PostgreSQL stream listener loop on channel", channel, ":", err)
|
|
time.Sleep(time.Second * 1)
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Special fixes for horizontally deployed pocketbase instances.
|
|
// When a user update settings or collection definitions, the changes is now only
|
|
// applied to the current instance where the dashboard is connected to.
|
|
// We need to broadcast the changes to all other instances.
|
|
// We leverage existing `shared_bridge_channel` notification channel to notify other instances.
|
|
func (t *RealtimeBridge) broadcastCollectionChanged(e *core.CollectionRequestEvent) error {
|
|
if err := e.Next(); err != nil {
|
|
return err
|
|
}
|
|
sql := `SELECT pg_notify('shared_bridge_channel', 'collection_updated|' || 'empty-payload')`
|
|
if _, err := t.app.DB().NewQuery(sql).Execute(); err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error sending [collection_updated] notification:", err)
|
|
// ignore the error as it is not critical
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Special fixes for horizontally deployed pocketbase instances.
|
|
// When a user update settings or collection definitions, the changes is now only
|
|
// applied to the current instance where the dashboard is connected to.
|
|
// We need to broadcast the changes to all other instances.
|
|
// We leverage existing `shared_bridge_channel` notification channel to notify other instances.
|
|
func (t *RealtimeBridge) broadcastSettingsUpdated(e *core.SettingsUpdateRequestEvent) error {
|
|
if err := e.Next(); err != nil {
|
|
return err
|
|
}
|
|
sql := `SELECT pg_notify('shared_bridge_channel', 'settings_updated|' || 'empty-payload')`
|
|
if _, err := t.app.DB().NewQuery(sql).Execute(); err != nil {
|
|
fmt.Fprintln(os.Stderr, "Error sending [settings_updated] notification:", err)
|
|
// ignore the error as it is not critical
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// runPgxCommand aquires a connection from connection pool, runs the provided pgxFunc,
|
|
// and put the connection back to the pool.
|
|
// *dbx.DB -> *sql.DB -> *sql.Conn -> *pgx.Conn
|
|
func runPgxCommand(ctx context.Context, app core.App, pgxFunc func(*pgx.Conn) error) error {
|
|
dbxDB, _ := app.NonconcurrentDB().(*dbx.DB)
|
|
if dbxDB == nil {
|
|
return fmt.Errorf("app.NonconcurrentDB() is not a *dbx.DB instance")
|
|
}
|
|
sqlDB := dbxDB.DB() // *sql.DB, a connection pool
|
|
sqlConn, err := sqlDB.Conn(ctx)
|
|
if err != nil {
|
|
return fmt.Errorf("error acquiring SQL connection: %w", err)
|
|
}
|
|
defer sqlConn.Close() // Put back to connection pool
|
|
|
|
return sqlConn.Raw(func(driverConn any) error {
|
|
stdlibConn, ok := driverConn.(*stdlib.Conn)
|
|
if !ok {
|
|
return fmt.Errorf("driverConn is not a *stdlib.Conn instance, are you using the pgx's database/sql driver?")
|
|
}
|
|
pgxConn := stdlibConn.Conn()
|
|
return pgxFunc(pgxConn)
|
|
})
|
|
}
|