commit 6cef528fe447cbe57f933ec6389f99bcdc7f395d · edavis.dev/monarch

+6 -65

cmd/monarch/cursors.go

···

       2
        
       

     

       3
        
       import (

     

       4
        
       	"context"

     

       5
       -
       	"fmt"

     

       6
        
       	"log/slog"

     

       7
       -
       	"strconv"

     

       8
        
       	"sync"

     

       9
        
       	"time"

     

       10
        
       

     

       11
        
       	"gorm.io/gorm"

     

       12
        
       )

     

       13
        
       

     

       14
       -
       type cursorRecord struct {

     

       15
       -
       	ID  uint `gorm:"primaryKey"`

     

       16
       -
       	Key string

     

       17
       -
       	Val string

     

       18
       -
       }

     

       19
       -
       

     

       20
        
       type CursorService struct {

     

       21
        
       	store *gorm.DB

     

       22
        
       

     

       23
        
       	firehoseLk  sync.Mutex

     

       24
        
       	firehoseSeq int64

     

       25
       -
       

     

       26
       -
       	reposLk  sync.Mutex

     

       27
       -
       	reposSeq string

     

       28
        
       }

     

       29
        
       

     

       30
        
       func NewCursorService(store *gorm.DB) *CursorService {

     

       31
       -
       	store.AutoMigrate(&cursorRecord{})

     

       32
       -
       

     

       33
       -
       	var rec cursorRecord

     

       34
       -
       	store.First(&rec, 1)

     

       35
       -
       	if rec.ID == 0 {

     

       36
       -
       		store.Create(&cursorRecord{ID: 1, Key: "firehose", Val: ""})

     

       37
       -
       	}

     

       38
       -
       

     

       39
       -
       	store.First(&rec, 2)

     

       40
       -
       	if rec.ID == 0 {

     

       41
       -
       		store.Create(&cursorRecord{ID: 2, Key: "repos", Val: ""})

     

       42
       -
       	}

     

       43
        
       

     

       44
        
       	return &CursorService{

     

       45
        
       		store: store,

     

       46
        
       	}

     

       47
        
       }

     

       48
        
       

     

       49
       -
       func (cs *CursorService) Get(key string) (string, error) {

     

       50
       -
       	var rec cursorRecord

     

       51
       -
       	if err := cs.store.Where("key = ?", key).First(&rec).Error; err != nil {

     

       52
       -
       		return "", fmt.Errorf("error fetching cursor record: %w", err)

     

       53
       -
       	}

     

       54
       -
       	return rec.Val, nil

     

       55
       -
       }

     

       56
       -
       

     

       57
       -
       func (cs *CursorService) SetFirehoseCursor(seq int64) {

     

       58
       -
       	cs.firehoseLk.Lock()

     

       59
       -
       	cs.firehoseSeq = seq

     

       60
       -
       	cs.firehoseLk.Unlock()

     

       61
       -
       }

     

       62
       -
       

     

       63
       -
       func (cs *CursorService) SetReposCursor(value string) {

     

       64
       -
       	cs.reposLk.Lock()

     

       65
       -
       	cs.reposSeq = value

     

       66
       -
       	cs.reposLk.Unlock()

     

       67
       -
       }

     

       68
       -
       

     

       69
       -
       func (cs *CursorService) Flush() error {

     

       70
       -
       	cs.firehoseLk.Lock()

     

       71
       -
       	fcursor := strconv.Itoa(int(cs.firehoseSeq))

     

       72
       -
       	if err := cs.store.Model(&cursorRecord{}).Where("key = ?", "firehose").Update("val", fcursor).Error; err != nil {

     

       73
       -
       		return fmt.Errorf("error updating cursor record: %w", err)

     

       74
       -
       	}

     

       75
       -
       	cs.firehoseLk.Unlock()

     

       76
       -
       

     

       77
       -
       	cs.reposLk.Lock()

     

       78
       -
       	if err := cs.store.Model(&cursorRecord{}).Where("key = ?", "repos").Update("val", cs.reposSeq).Error; err != nil {

     

       79
       -
       		return fmt.Errorf("error updating cursor record: %w", err)

     

       80
       -
       	}

     

       81
       -
       	cs.reposLk.Unlock()

     

       82
       -
       

     

       83
       -
       	return nil

     

       84
       -
       }

     

       85
       -
       

     

       86
       -
       func (cs *CursorService) CheckpointCursors(ctx context.Context) {

     

       87
        
       	t := time.NewTicker(time.Second * 5)

     

       88
        
       	defer t.Stop()

     

       89
        
       

     

       90
        
       	for {

     

       91
        
       		select {

     

       92
        
       		case <-ctx.Done():

     

       93
       -
       			slog.Info("stopping cursor checkpointer")

     

       94
        
       			return

     

       95
        
       		case <-t.C:

     

       96
        
       		}

     

       97
        
       

     

       98
       -
       		slog.Info("flushing cursors", "firehose", cs.firehoseSeq, "repos", cs.reposSeq)

     

       99
       -
       		if err := cs.Flush(); err != nil {

     

       100
       -
       			slog.Error("error flushing cursors", "err", err)

     

       101
        
       			return

     

       102
        
       		}

     

       103
        
       	}

···

       2
        
       

     

       3
        
       import (

     

       4
        
       	"context"

     

       0
        
       
     

       5
        
       	"log/slog"

     

       0
        
       
     

       6
        
       	"sync"

     

       7
        
       	"time"

     

       8
        
       

     

       9
        
       	"gorm.io/gorm"

     

       10
        
       )

     

       11
        
       

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       12
        
       type CursorService struct {

     

       13
        
       	store *gorm.DB

     

       14
        
       

     

       15
        
       	firehoseLk  sync.Mutex

     

       16
        
       	firehoseSeq int64

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       17
        
       }

     

       18
        
       

     

       19
        
       func NewCursorService(store *gorm.DB) *CursorService {

     

       20
       +
       	store.AutoMigrate(&firehoseCursor{})

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       21
        
       

     

       22
        
       	return &CursorService{

     

       23
        
       		store: store,

     

       24
        
       	}

     

       25
        
       }

     

       26
        
       

     

       27
       +
       func (cs *CursorService) Checkpoint(ctx context.Context) {

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       28
        
       	t := time.NewTicker(time.Second * 5)

     

       29
        
       	defer t.Stop()

     

       30
        
       

     

       31
        
       	for {

     

       32
        
       		select {

     

       33
        
       		case <-ctx.Done():

     

       34
       +
       			slog.Info("stopping cursor checkpointer", "err", ctx.Err())

     

       35
        
       			return

     

       36
        
       		case <-t.C:

     

       37
        
       		}

     

       38
        
       

     

       39
       +
       		slog.Info("persisting firehose cursor", "seq", cs.firehoseSeq)

     

       40
       +
       		if err := cs.PersistFirehoseCursor(); err != nil {

     

       41
       +
       			slog.Error("error persisting firehose cursor", "err", err)

     

       42
        
       			return

     

       43
        
       		}

     

       44
        
       	}

+38 -3

cmd/monarch/firehose.go

···

       11
        
       

     

       12
        
       func NewFirehoseConnection(ctx context.Context, cctx *cli.Context, cursorSvc *CursorService) (*websocket.Conn, error) {

     

       13
        
       	url := fmt.Sprintf("wss://%s/xrpc/com.atproto.sync.subscribeRepos", cctx.String("relay-host"))

     

       14
       -
       	curs, _ := cursorSvc.Get("firehose")

     

       15
       -
       	if curs != "" {

     

       16
       -
       		url += "?cursor=" + curs

     

       17
        
       	}

     

       18
        
       

     

       19
        
       	conn, _, err := websocket.DefaultDialer.DialContext(ctx, url, http.Header{

     
···

       25
        
       

     

       26
        
       	return conn, nil

     

       27
        
       }

     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0
        
       
     

       0

···

       11
        
       

     

       12
        
       func NewFirehoseConnection(ctx context.Context, cctx *cli.Context, cursorSvc *CursorService) (*websocket.Conn, error) {

     

       13
        
       	url := fmt.Sprintf("wss://%s/xrpc/com.atproto.sync.subscribeRepos", cctx.String("relay-host"))

     

       14
       +
       	curs, err := cursorSvc.GetFirehoseCursor()

     

       15
       +
       	if err == nil { // reversed

     

       16
       +
       		url += fmt.Sprintf("?cursor=%d", curs)

     

       17
        
       	}

     

       18
        
       

     

       19
        
       	conn, _, err := websocket.DefaultDialer.DialContext(ctx, url, http.Header{

     
···

       25
        
       

     

       26
        
       	return conn, nil

     

       27
        
       }

     

       28
       +
       

     

       29
       +
       type firehoseCursor struct {

     

       30
       +
       	ID int `gorm:"primaryKey"`

     

       31
       +
       	Key string

     

       32
       +
       	Val int64

     

       33
       +
       }

     

       34
       +
       

     

       35
       +
       func (cs *CursorService) GetFirehoseCursor() (int64, error) {

     

       36
       +
       	cs.firehoseLk.Lock()

     

       37
       +
       	defer cs.firehoseLk.Unlock()

     

       38
       +
       

     

       39
       +
       	var fcur firehoseCursor

     

       40
       +
       	if err := cs.store.Where("key = ?", "firehose").Attrs(firehoseCursor{Val: cs.firehoseSeq}).FirstOrCreate(&fcur).Error; err != nil {

     

       41
       +
       		return 0, fmt.Errorf("error getting firehose seq from DB: %w", err)

     

       42
       +
       	}

     

       43
       +
       	return fcur.Val, nil

     

       44
       +
       }

     

       45
       +
       

     

       46
       +
       func (cs *CursorService) SetFirehoseCursor(seq int64) {

     

       47
       +
       	cs.firehoseLk.Lock()

     

       48
       +
       	defer cs.firehoseLk.Unlock()

     

       49
       +
       

     

       50
       +
       	cs.firehoseSeq = seq

     

       51
       +
       }

     

       52
       +
       

     

       53
       +
       func (cs *CursorService) PersistFirehoseCursor() error {

     

       54
       +
       	cs.firehoseLk.Lock()

     

       55
       +
       	defer cs.firehoseLk.Unlock()

     

       56
       +
       

     

       57
       +
       	var fcur firehoseCursor

     

       58
       +
       	if err := cs.store.Where("key = ?", "firehose").Assign(firehoseCursor{Val: cs.firehoseSeq}).FirstOrCreate(&fcur).Error; err != nil {

     

       59
       +
       		return fmt.Errorf("error persisting firehose seq: %w", err)

     

       60
       +
       	}

     

       61
       +
       	return nil

     

       62
       +
       }

+1 -1

cmd/monarch/main.go

···

       42
        
       	slog.Info("starting up")

     

       43
        
       

     

       44
        
       	app.cursor = NewCursorService(app.state)

     

       45
       -
       	go app.cursor.CheckpointCursors(ctx)

     

       46
        
       

     

       47
        
       	app.handler = NewHandlerService(app.content)

     

       48

···

       42
        
       	slog.Info("starting up")

     

       43
        
       

     

       44
        
       	app.cursor = NewCursorService(app.state)

     

       45
       +
       	go app.cursor.Checkpoint(ctx)

     

       46
        
       

     

       47
        
       	app.handler = NewHandlerService(app.content)

     

       48