commit cfbcc31012529bcaec101039e1b70668538002b0 · hailey.at/cocoon

+31 -11

server/handle_sync_subscribe_repos.go

···

       1
       1
        
       package server

     

       2
       2
        
       

     

       3
       3
        
       import (

     

       4
       4
       -
       	"fmt"

     

       4
       4
       +
       	"context"

     

       5
       5
       +
       	"time"

     

       5
       6
        
       

     

       6
       7
        
       	"github.com/bluesky-social/indigo/events"

     

       7
       8
        
       	"github.com/bluesky-social/indigo/lex/util"

     
···

       10
       11
        
       )

     

       11
       12
        
       

     

       12
       13
        
       func (s *Server) handleSyncSubscribeRepos(e echo.Context) error {

     

       14
       14
       +
       	ctx := e.Request().Context()

     

       15
       15
       +
       	logger := s.logger.With("component", "subscribe-repos-websocket")

     

       16
       16
       +
       

     

       13
       17
        
       	conn, err := websocket.Upgrade(e.Response().Writer, e.Request(), e.Response().Header(), 1<<10, 1<<10)

     

       14
       18
        
       	if err != nil {

     

       19
       19
       +
       		logger.Error("unable to establish websocket with relay", "err", err)

     

       15
       20
        
       		return err

     

       16
       21
        
       	}

     

       17
       17
       -
       

     

       18
       18
       -
       	s.logger.Info("new connection", "ua", e.Request().UserAgent())

     

       19
       19
       -
       

     

       20
       20
       -
       	ctx := e.Request().Context()

     

       21
       22
        
       

     

       22
       23
        
       	ident := e.RealIP() + "-" + e.Request().UserAgent()

     

       24
       24
       +
       	logger = logger.With("ident", ident)

     

       25
       25
       +
       	logger.Info("new connection established")

     

       23
       26
        
       

     

       24
       27
        
       	evts, cancel, err := s.evtman.Subscribe(ctx, ident, func(evt *events.XRPCStreamEvent) bool {

     

       25
       28
        
       		return true

     
···

       33
       36
        
       	for evt := range evts {

     

       34
       37
        
       		wc, err := conn.NextWriter(websocket.BinaryMessage)

     

       35
       38
        
       		if err != nil {

     

       36
       36
       -
       			return err

     

       39
       39
       +
       			logger.Error("error writing message to relay", "err", err)

     

       40
       40
       +
       			break

     

       37
       41
        
       		}

     

       38
       42
        
       

     

       39
       39
       -
       		var obj util.CBOR

     

       43
       43
       +
       		if ctx.Err() != nil {

     

       44
       44
       +
       			logger.Error("context error", "err", err)

     

       45
       45
       +
       			break

     

       46
       46
       +
       		}

     

       40
       47
        
       

     

       48
       48
       +
       		var obj util.CBOR

     

       41
       49
        
       		switch {

     

       42
       50
        
       		case evt.Error != nil:

     

       43
       51
        
       			header.Op = events.EvtKindErrorFrame

     
···

       55
       63
        
       			header.MsgType = "#info"

     

       56
       64
        
       			obj = evt.RepoInfo

     

       57
       65
        
       		default:

     

       58
       58
       -
       			return fmt.Errorf("unrecognized event kind")

     

       66
       66
       +
       			logger.Warn("unrecognized event kind")

     

       67
       67
       +
       			return nil

     

       59
       68
        
       		}

     

       60
       69
        
       

     

       61
       70
        
       		if err := header.MarshalCBOR(wc); err != nil {

     

       62
       62
       -
       			return fmt.Errorf("failed to write header: %w", err)

     

       71
       71
       +
       			logger.Error("failed to write header to relay", "err", err)

     

       72
       72
       +
       			break

     

       63
       73
        
       		}

     

       64
       74
        
       

     

       65
       75
        
       		if err := obj.MarshalCBOR(wc); err != nil {

     

       66
       66
       -
       			return fmt.Errorf("failed to write event: %w", err)

     

       76
       76
       +
       			logger.Error("failed to write event to relay", "err", err)

     

       77
       77
       +
       			break

     

       67
       78
        
       		}

     

       68
       79
        
       

     

       69
       80
        
       		if err := wc.Close(); err != nil {

     

       70
       70
       -
       			return fmt.Errorf("failed to flush-close our event write: %w", err)

     

       81
       81
       +
       			logger.Error("failed to flush-close our event write", "err", err)

     

       82
       82
       +
       			break

     

       71
       83
        
       		}

     

       84
       84
       +
       	}

     

       85
       85
       +
       

     

       86
       86
       +
       	// we should tell the relay to request a new crawl at this point if we got disconnected

     

       87
       87
       +
       	// use a new context since the old one might be cancelled at this point

     

       88
       88
       +
       	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)

     

       89
       89
       +
       	defer cancel()

     

       90
       90
       +
       	if err := s.requestCrawl(ctx); err != nil {

     

       91
       91
       +
       		logger.Error("error requesting crawls", "err", err)

     

       72
       92
        
       	}

     

       73
       93
        
       

     

       74
       94
        
       	return nil

+36 -5

server/server.go

···

       77
       77
        
       	passport      *identity.Passport

     

       78
       78
        
       	fallbackProxy string

     

       79
       79
        
       

     

       80
       80
       +
       	lastRequestCrawl time.Time

     

       81
       81
       +
       	requestCrawlMu   sync.Mutex

     

       82
       82
       +
       

     

       80
       83
        
       	dbName   string

     

       81
       84
        
       	s3Config *S3Config

     

       82
       85
        
       }

     
···

       518
       521
        
       

     

       519
       522
        
       	go s.backupRoutine()

     

       520
       523
        
       

     

       524
       524
       +
       	go func() {

     

       525
       525
       +
       		if err := s.requestCrawl(ctx); err != nil {

     

       526
       526
       +
       			s.logger.Error("error requesting crawls", "err", err)

     

       527
       527
       +
       		}

     

       528
       528
       +
       	}()

     

       529
       529
       +
       

     

       530
       530
       +
       	<-ctx.Done()

     

       531
       531
       +
       

     

       532
       532
       +
       	fmt.Println("shut down")

     

       533
       533
       +
       

     

       534
       534
       +
       	return nil

     

       535
       535
       +
       }

     

       536
       536
       +
       

     

       537
       537
       +
       func (s *Server) requestCrawl(ctx context.Context) error {

     

       538
       538
       +
       	logger := s.logger.With("component", "request-crawl")

     

       539
       539
       +
       	s.requestCrawlMu.Lock()

     

       540
       540
       +
       	defer s.requestCrawlMu.Unlock()

     

       541
       541
       +
       

     

       542
       542
       +
       	logger.Info("requesting crawl with configured relays")

     

       543
       543
       +
       

     

       544
       544
       +
       	if time.Now().Sub(s.lastRequestCrawl) <= 1*time.Minute {

     

       545
       545
       +
       		return fmt.Errorf("a crawl request has already been made within the last minute")

     

       546
       546
       +
       	}

     

       547
       547
       +
       

     

       521
       548
        
       	for _, relay := range s.config.Relays {

     

       549
       549
       +
       		logger := logger.With("relay", relay)

     

       550
       550
       +
       		logger.Info("requesting crawl from relay")

     

       522
       551
        
       		cli := xrpc.Client{Host: relay}

     

       523
       523
       -
       		atproto.SyncRequestCrawl(ctx, &cli, &atproto.SyncRequestCrawl_Input{

     

       552
       552
       +
       		if err := atproto.SyncRequestCrawl(ctx, &cli, &atproto.SyncRequestCrawl_Input{

     

       524
       553
        
       			Hostname: s.config.Hostname,

     

       525
       525
       -
       		})

     

       554
       554
       +
       		}); err != nil {

     

       555
       555
       +
       			logger.Error("error requesting crawl", "err", err)

     

       556
       556
       +
       		} else {

     

       557
       557
       +
       			logger.Info("crawl requested successfully")

     

       558
       558
       +
       		}

     

       526
       559
        
       	}

     

       527
       560
        
       

     

       528
       528
       -
       	<-ctx.Done()

     

       529
       529
       -
       

     

       530
       530
       -
       	fmt.Println("shut down")

     

       561
       561
       +
       	s.lastRequestCrawl = time.Now()

     

       531
       562
        
       

     

       532
       563
        
       	return nil

     

       533
       564
        
       }