Something went wrong. Try again.
A fork of the Cocoon PDS but being made more distributed. forked from willdot.net/cocoon
Something went wrong. Try again.
5.8 kB · 252 lines
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253package server
import ( "bytes" "context" "errors" "fmt" "log/slog" "sync" "time"
"github.com/avast/retry-go/v4" "github.com/bluesky-social/indigo/api/atproto" "github.com/bluesky-social/indigo/events" indigomodels "github.com/bluesky-social/indigo/models" cbg "github.com/whyrusleeping/cbor-gen" "gorm.io/gorm"
"github.com/haileyok/cocoon/models")
type DbPersister struct { Db *gorm.DB
Lk sync.Mutex
Broadcast func(*events.XRPCStreamEvent)
// how long do we actually want to keep these things around Retention time.Duration}
func NewDbPersister(db *gorm.DB, retention time.Duration) (*DbPersister, error) { if err := db.AutoMigrate(&models.EventRecord{}); err != nil { return nil, fmt.Errorf("failed to migrate EventRecord: %w", err) }
if retention == 0 { retention = 72 * time.Hour }
p := &DbPersister{ Db: db, Retention: retention, }
go p.cleanupRoutine()
return p, nil}
func (p *DbPersister) SetEventBroadcaster(brc func(*events.XRPCStreamEvent)) { p.Broadcast = brc}
func (p *DbPersister) Persist(ctx context.Context, e *events.XRPCStreamEvent) error { p.Lk.Lock() defer p.Lk.Unlock()
rec := &models.EventRecord{} if err := p.Db.Order("seq desc").Limit(1).First(rec).Error; err != nil { slog.Error("fetching most recent event record", "error", err) rec.Seq = time.Now().UnixMilli() }
// if the error on inserting the event record is a constraint error, it means that // another event has been stored since getting the last sequence number. In that case // retry which will increase the sequence number again and hopefully work. Any others // can error out. retryIfFunc := func(err error) bool { return errors.Is(err, gorm.ErrDuplicatedKey) } err := retry.Do(func() error { rec.Seq++
var did string var evtType string
switch { case e.RepoCommit != nil: e.RepoCommit.Seq = rec.Seq did = e.RepoCommit.Repo evtType = "commit" case e.RepoSync != nil: e.RepoSync.Seq = rec.Seq did = e.RepoSync.Did evtType = "sync" case e.RepoIdentity != nil: e.RepoIdentity.Seq = rec.Seq did = e.RepoIdentity.Did evtType = "identity" case e.RepoAccount != nil: e.RepoAccount.Seq = rec.Seq did = e.RepoAccount.Did evtType = "account" default: return fmt.Errorf("unknown event type") }
data, err := serializeEvent(e) if err != nil { return fmt.Errorf("failed to serialize event: %w", err) }
rec.CreatedAt = time.Now() rec.Did = did rec.Type = evtType rec.Data = data
err = p.Db.Create(rec).Error if err != nil { return fmt.Errorf("failed to persist event: %w", err) }
return nil }, retry.RetryIf(retryIfFunc))
if err != nil { return err }
if p.Broadcast != nil { p.Broadcast(e) }
return nil}
func (p *DbPersister) Playback(ctx context.Context, since int64, cb func(*events.XRPCStreamEvent) error) error { const pageSize = 500
cursor := since for { var records []models.EventRecord if err := p.Db.WithContext(ctx). Where("seq > ?", cursor). Order("seq asc"). Limit(pageSize). Find(&records).Error; err != nil { return fmt.Errorf("failed to query events: %w", err) }
slog.Info("playback", "len of records", len(records), "cursor", cursor)
if len(records) == 0 { return nil }
for _, rec := range records { evt, err := deserializeEvent(rec.Type, rec.Data) if err != nil { return fmt.Errorf("failed to deserialize event %d: %w", rec.Seq, err) }
if err := cb(evt); err != nil { return err }
cursor = rec.Seq }
if len(records) < pageSize { return nil } }}
func (p *DbPersister) TakeDownRepo(ctx context.Context, uid indigomodels.Uid) error { return nil}
func (p *DbPersister) Flush(ctx context.Context) error { return nil}
func (p *DbPersister) Shutdown(ctx context.Context) error { return nil}
func (p *DbPersister) cleanupRoutine() { ticker := time.NewTicker(time.Hour) defer ticker.Stop()
for range ticker.C { cutoff := time.Now().Add(-p.Retention) if err := p.Db.Where("created_at < ?", cutoff).Delete(&models.EventRecord{}).Error; err != nil { continue } }}
func serializeEvent(e *events.XRPCStreamEvent) ([]byte, error) { buf := new(bytes.Buffer) cw := cbg.NewCborWriter(buf)
switch { case e.RepoCommit != nil: if err := e.RepoCommit.MarshalCBOR(cw); err != nil { return nil, err } case e.RepoSync != nil: if err := e.RepoSync.MarshalCBOR(cw); err != nil { return nil, err } case e.RepoIdentity != nil: if err := e.RepoIdentity.MarshalCBOR(cw); err != nil { return nil, err } case e.RepoAccount != nil: if err := e.RepoAccount.MarshalCBOR(cw); err != nil { return nil, err } default: return nil, fmt.Errorf("unknown event type") }
return buf.Bytes(), nil}
func deserializeEvent(evtType string, data []byte) (*events.XRPCStreamEvent, error) { r := bytes.NewReader(data) cr := cbg.NewCborReader(r)
switch evtType { case "commit": evt := &atproto.SyncSubscribeRepos_Commit{} if err := evt.UnmarshalCBOR(cr); err != nil { return nil, err } return &events.XRPCStreamEvent{RepoCommit: evt}, nil case "sync": evt := &atproto.SyncSubscribeRepos_Sync{} if err := evt.UnmarshalCBOR(cr); err != nil { return nil, err } return &events.XRPCStreamEvent{RepoSync: evt}, nil case "identity": evt := &atproto.SyncSubscribeRepos_Identity{} if err := evt.UnmarshalCBOR(cr); err != nil { return nil, err } return &events.XRPCStreamEvent{RepoIdentity: evt}, nil case "account": evt := &atproto.SyncSubscribeRepos_Account{} if err := evt.UnmarshalCBOR(cr); err != nil { return nil, err } return &events.XRPCStreamEvent{RepoAccount: evt}, nil default: return nil, fmt.Errorf("unknown event type: %s", evtType) }}