Import
This commit is contained in:
Max Ignatenko
parent
2b6abac607
commit
63a767d890
25 changed files with 3027 additions and 0 deletions
14
cmd/lister/Dockerfile
Normal file
14
cmd/lister/Dockerfile
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
FROM golang:1.21.1 as builder
|
||||
WORKDIR /app
|
||||
COPY go.mod go.sum ./
|
||||
RUN go mod download
|
||||
COPY . ./
|
||||
RUN go build -trimpath ./cmd/lister
|
||||
|
||||
FROM alpine:latest as certs
|
||||
RUN apk --update add ca-certificates
|
||||
|
||||
FROM debian:stable-slim
|
||||
COPY --from=certs /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
|
||||
COPY --from=builder /app/lister .
|
||||
ENTRYPOINT ["./lister"]
|
||||
108
cmd/lister/lister.go
Normal file
108
cmd/lister/lister.go
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"github.com/rs/zerolog"
|
||||
"gorm.io/gorm"
|
||||
|
||||
comatproto "github.com/bluesky-social/indigo/api/atproto"
|
||||
"github.com/bluesky-social/indigo/did"
|
||||
|
||||
"github.com/uabluerail/bsky-tools/pagination"
|
||||
"github.com/uabluerail/bsky-tools/xrpcauth"
|
||||
"github.com/uabluerail/indexer/pds"
|
||||
"github.com/uabluerail/indexer/repo"
|
||||
"github.com/uabluerail/indexer/util/resolver"
|
||||
)
|
||||
|
||||
type Lister struct {
|
||||
db *gorm.DB
|
||||
resolver did.Resolver
|
||||
|
||||
pollInterval time.Duration
|
||||
listRefreshInterval time.Duration
|
||||
}
|
||||
|
||||
func NewLister(ctx context.Context, db *gorm.DB) (*Lister, error) {
|
||||
return &Lister{
|
||||
db: db,
|
||||
resolver: resolver.Resolver,
|
||||
pollInterval: 5 * time.Minute,
|
||||
listRefreshInterval: 24 * time.Hour,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (l *Lister) Start(ctx context.Context) error {
|
||||
go l.run(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *Lister) run(ctx context.Context) {
|
||||
log := zerolog.Ctx(ctx)
|
||||
ticker := time.NewTicker(l.pollInterval)
|
||||
|
||||
log.Info().Msgf("Lister starting...")
|
||||
t := make(chan time.Time, 1)
|
||||
t <- time.Now()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Info().Msgf("Lister stopped (context expired)")
|
||||
return
|
||||
case <-t:
|
||||
db := l.db.WithContext(ctx)
|
||||
|
||||
remote := pds.PDS{}
|
||||
if err := db.Model(&remote).
|
||||
Where("last_list is null or last_list < ?", time.Now().Add(-l.listRefreshInterval)).
|
||||
Take(&remote).Error; err != nil {
|
||||
if !errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
log.Error().Err(err).Msgf("Failed to query DB for a PDS to list repos from: %s", err)
|
||||
}
|
||||
break
|
||||
}
|
||||
client := xrpcauth.NewAnonymousClient(ctx)
|
||||
client.Host = remote.Host
|
||||
|
||||
log.Info().Msgf("Listing repos from %q...", remote.Host)
|
||||
dids, err := pagination.Reduce(
|
||||
func(cursor string) (resp *comatproto.SyncListRepos_Output, nextCursor string, err error) {
|
||||
resp, err = comatproto.SyncListRepos(ctx, client, cursor, 200)
|
||||
if err == nil && resp.Cursor != nil {
|
||||
nextCursor = *resp.Cursor
|
||||
}
|
||||
return
|
||||
},
|
||||
func(resp *comatproto.SyncListRepos_Output, acc []string) ([]string, error) {
|
||||
for _, repo := range resp.Repos {
|
||||
if repo == nil {
|
||||
continue
|
||||
}
|
||||
acc = append(acc, repo.Did)
|
||||
}
|
||||
return acc, nil
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
log.Error().Err(err).Msgf("Failed to list repos from %q: %s", remote.Host, err)
|
||||
break
|
||||
}
|
||||
log.Info().Msgf("Received %d DIDs from %q", len(dids), remote.Host)
|
||||
|
||||
for _, did := range dids {
|
||||
if _, err := repo.EnsureExists(ctx, l.db, did); err != nil {
|
||||
log.Error().Err(err).Msgf("Failed to ensure that we have a record for the repo %q: %s", did, err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := db.Model(&remote).Updates(&pds.PDS{LastList: time.Now()}).Error; err != nil {
|
||||
log.Error().Err(err).Msgf("Failed to update the timestamp of last list for %q: %s", remote.Host, err)
|
||||
}
|
||||
case v := <-ticker.C:
|
||||
t <- v
|
||||
}
|
||||
}
|
||||
}
|
||||
180
cmd/lister/main.go
Normal file
180
cmd/lister/main.go
Normal file
|
|
@ -0,0 +1,180 @@
|
|||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
_ "net/http/pprof"
|
||||
"os"
|
||||
"os/signal"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"runtime/debug"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
_ "github.com/joho/godotenv/autoload"
|
||||
"github.com/kelseyhightower/envconfig"
|
||||
"github.com/prometheus/client_golang/prometheus/promhttp"
|
||||
"github.com/rs/zerolog"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"github.com/uabluerail/indexer/pds"
|
||||
"github.com/uabluerail/indexer/repo"
|
||||
"github.com/uabluerail/indexer/util/gormzerolog"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
LogFile string
|
||||
LogFormat string `default:"text"`
|
||||
LogLevel int64 `default:"1"`
|
||||
MetricsPort string `split_words:"true"`
|
||||
DBUrl string `envconfig:"POSTGRES_URL"`
|
||||
}
|
||||
|
||||
var config Config
|
||||
|
||||
func runMain(ctx context.Context) error {
|
||||
ctx = setupLogging(ctx)
|
||||
log := zerolog.Ctx(ctx)
|
||||
log.Debug().Msgf("Starting up...")
|
||||
db, err := gorm.Open(postgres.Open(config.DBUrl), &gorm.Config{
|
||||
Logger: gormzerolog.New(&logger.Config{
|
||||
SlowThreshold: 1 * time.Second,
|
||||
IgnoreRecordNotFoundError: true,
|
||||
}, nil),
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("connecting to the database: %w", err)
|
||||
}
|
||||
log.Debug().Msgf("DB connection established")
|
||||
|
||||
for _, f := range []func(*gorm.DB) error{
|
||||
pds.AutoMigrate,
|
||||
repo.AutoMigrate,
|
||||
} {
|
||||
if err := f(db); err != nil {
|
||||
return fmt.Errorf("auto-migrating DB schema: %w", err)
|
||||
}
|
||||
}
|
||||
log.Debug().Msgf("DB schema updated")
|
||||
|
||||
lister, err := NewLister(ctx, db)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create lister: %w", err)
|
||||
}
|
||||
if err := lister.Start(ctx); err != nil {
|
||||
return fmt.Errorf("failed to start lister: %w", err)
|
||||
}
|
||||
|
||||
log.Info().Msgf("Starting HTTP listener on %q...", config.MetricsPort)
|
||||
http.Handle("/metrics", promhttp.Handler())
|
||||
srv := &http.Server{Addr: fmt.Sprintf(":%s", config.MetricsPort)}
|
||||
errCh := make(chan error)
|
||||
go func() {
|
||||
errCh <- srv.ListenAndServe()
|
||||
}()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
if err := srv.Shutdown(context.Background()); err != nil {
|
||||
return fmt.Errorf("HTTP server shutdown failed: %w", err)
|
||||
}
|
||||
}
|
||||
return <-errCh
|
||||
}
|
||||
|
||||
func main() {
|
||||
flag.StringVar(&config.LogFile, "log", "", "Path to the log file. If empty, will log to stderr")
|
||||
flag.StringVar(&config.LogFormat, "log-format", "text", "Logging format. 'text' or 'json'")
|
||||
flag.Int64Var(&config.LogLevel, "log-level", 1, "Log level. -1 - trace, 0 - debug, 1 - info, 5 - panic")
|
||||
|
||||
if err := envconfig.Process("lister", &config); err != nil {
|
||||
log.Fatalf("envconfig.Process: %s", err)
|
||||
}
|
||||
|
||||
flag.Parse()
|
||||
|
||||
ctx, _ := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
if err := runMain(ctx); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func setupLogging(ctx context.Context) context.Context {
|
||||
logFile := os.Stderr
|
||||
|
||||
if config.LogFile != "" {
|
||||
f, err := os.OpenFile(config.LogFile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
log.Fatalf("Failed to open the specified log file %q: %s", config.LogFile, err)
|
||||
}
|
||||
logFile = f
|
||||
}
|
||||
|
||||
var output io.Writer
|
||||
|
||||
switch config.LogFormat {
|
||||
case "json":
|
||||
output = logFile
|
||||
case "text":
|
||||
prefixList := []string{}
|
||||
info, ok := debug.ReadBuildInfo()
|
||||
if ok {
|
||||
prefixList = append(prefixList, info.Path+"/")
|
||||
}
|
||||
|
||||
basedir := ""
|
||||
_, sourceFile, _, ok := runtime.Caller(0)
|
||||
if ok {
|
||||
basedir = filepath.Dir(sourceFile)
|
||||
}
|
||||
|
||||
if basedir != "" && strings.HasPrefix(basedir, "/") {
|
||||
prefixList = append(prefixList, basedir+"/")
|
||||
head, _ := filepath.Split(basedir)
|
||||
for head != "/" {
|
||||
prefixList = append(prefixList, head)
|
||||
head, _ = filepath.Split(strings.TrimSuffix(head, "/"))
|
||||
}
|
||||
}
|
||||
|
||||
output = zerolog.ConsoleWriter{
|
||||
Out: logFile,
|
||||
NoColor: true,
|
||||
TimeFormat: time.RFC3339,
|
||||
PartsOrder: []string{
|
||||
zerolog.LevelFieldName,
|
||||
zerolog.TimestampFieldName,
|
||||
zerolog.CallerFieldName,
|
||||
zerolog.MessageFieldName,
|
||||
},
|
||||
FormatFieldName: func(i interface{}) string { return fmt.Sprintf("%s:", i) },
|
||||
FormatFieldValue: func(i interface{}) string { return fmt.Sprintf("%s", i) },
|
||||
FormatCaller: func(i interface{}) string {
|
||||
s := i.(string)
|
||||
for _, p := range prefixList {
|
||||
s = strings.TrimPrefix(s, p)
|
||||
}
|
||||
return s
|
||||
},
|
||||
}
|
||||
default:
|
||||
log.Fatalf("Invalid log format specified: %q", config.LogFormat)
|
||||
}
|
||||
|
||||
logger := zerolog.New(output).Level(zerolog.Level(config.LogLevel)).With().Caller().Timestamp().Logger()
|
||||
|
||||
ctx = logger.WithContext(ctx)
|
||||
|
||||
zerolog.DefaultContextLogger = &logger
|
||||
log.SetOutput(logger)
|
||||
|
||||
return ctx
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue