apollo-backend/cmd/apollo-scheduler/main.go

285 lines
6.2 KiB
Go
Raw Normal View History

2021-07-08 23:03:46 +00:00
package main
import (
"context"
2021-07-09 03:12:50 +00:00
"encoding/json"
2021-07-08 23:03:46 +00:00
"fmt"
"log"
"os"
"os/signal"
"strconv"
2021-07-08 23:03:46 +00:00
"syscall"
"time"
2021-07-09 02:09:14 +00:00
"github.com/DataDog/datadog-go/statsd"
2021-07-08 23:03:46 +00:00
"github.com/adjust/rmq/v4"
"github.com/go-co-op/gocron"
"github.com/go-redis/redis/v8"
"github.com/jackc/pgx/v4"
"github.com/jackc/pgx/v4/pgxpool"
"github.com/joho/godotenv"
"github.com/sirupsen/logrus"
)
2021-07-09 03:12:50 +00:00
const (
2021-07-09 04:27:50 +00:00
batchSize = 100
checkTimeout = 60 // how long until we force a check
enqueueTimeout = 5 // how long until we try to re-enqueue
2021-07-09 03:12:50 +00:00
)
2021-07-08 23:03:46 +00:00
func main() {
_ = godotenv.Load()
errChan := make(chan error, 10)
go logErrors(errChan)
ctx, cancel := context.WithCancel(context.Background())
var logger *logrus.Logger
{
logger = logrus.New()
if os.Getenv("ENV") == "" {
logger.SetLevel(logrus.DebugLevel)
} else {
logger.SetFormatter(&logrus.TextFormatter{
DisableColors: true,
FullTimestamp: true,
})
}
}
2021-07-09 02:09:14 +00:00
statsd, err := statsd.New("127.0.0.1:8125")
if err != nil {
logger.WithFields(logrus.Fields{
"err": err,
}).Error("failed to set up stats")
}
2021-07-08 23:03:46 +00:00
// Set up Postgres connection
var pool *pgxpool.Pool
{
config, err := pgxpool.ParseConfig(os.Getenv("DATABASE_CONNECTION_POOL_URL"))
if err != nil {
panic(err)
}
// Setting the build statement cache to nil helps this work with pgbouncer
config.ConnConfig.BuildStatementCache = nil
config.ConnConfig.PreferSimpleProtocol = true
pool, err = pgxpool.ConnectConfig(ctx, config)
if err != nil {
panic(err)
}
defer pool.Close()
}
// Set up Redis connection
var redisConn *redis.Client
{
2021-07-08 23:26:15 +00:00
opt, err := redis.ParseURL(os.Getenv("REDISCLOUD_URL"))
if err != nil {
panic(err)
}
2021-07-08 23:03:46 +00:00
2021-07-08 23:26:15 +00:00
redisConn = redis.NewClient(opt)
2021-07-08 23:03:46 +00:00
if err := redisConn.Ping(ctx).Err(); err != nil {
panic(err)
}
}
// Set up queues
var (
notificationsQueue rmq.Queue
)
{
connection, err := rmq.OpenConnectionWithRedisClient("producer", redisConn, errChan)
if err != nil {
panic(err)
}
notificationsQueue, err = connection.OpenQueue("notifications")
if err != nil {
panic(err)
}
}
s := gocron.NewScheduler(time.UTC)
2021-07-09 02:09:14 +00:00
s.Every(1).Second().Do(func() { enqueueAccounts(ctx, logger, statsd, pool, redisConn, notificationsQueue) })
s.Every(1).Minute().Do(func() { reportStats(ctx, logger, statsd, pool, redisConn) })
2021-07-08 23:03:46 +00:00
s.StartAsync()
signals := make(chan os.Signal, 1)
2021-07-09 01:07:01 +00:00
signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
2021-07-08 23:03:46 +00:00
defer signal.Stop(signals)
<-signals // wait for signal
cancel()
go func() {
<-signals // hard exit on second signal (in case shutdown gets stuck)
os.Exit(1)
}()
s.Stop()
}
func reportStats(ctx context.Context, logger *logrus.Logger, statsd *statsd.Client, pool *pgxpool.Pool, redisConn *redis.Client) {
var (
count int64
metrics = []struct {
query string
name string
}{
{"SELECT COUNT(*) FROM accounts", "apollo.registrations.accounts"},
{"SELECT COUNT(*) FROM devices", "apollo.registrations.devices"},
}
)
for _, metric := range metrics {
pool.QueryRow(ctx, metric.query).Scan(&count)
statsd.Gauge(metric.name, float64(count), []string{}, 1)
logger.WithFields(logrus.Fields{
"count": count,
"metric": metric.name,
}).Debug("fetched metrics")
}
}
2021-07-09 02:09:14 +00:00
func enqueueAccounts(ctx context.Context, logger *logrus.Logger, statsd *statsd.Client, pool *pgxpool.Pool, redisConn *redis.Client, queue rmq.Queue) {
2021-07-09 02:15:28 +00:00
start := time.Now()
2021-07-09 01:07:01 +00:00
now := float64(time.Now().UnixNano()/int64(time.Millisecond)) / 1000
2021-07-09 02:09:14 +00:00
// Start looking for accounts that were last checked at least 5 seconds ago
// and at most 6 seconds ago. Also look for accounts that haven't been checked
// in over a minute.
2021-07-09 02:15:28 +00:00
ts := start.Unix()
2021-07-09 04:27:50 +00:00
ready := ts - enqueueTimeout
expired := ts - checkTimeout
2021-07-09 00:53:09 +00:00
2021-07-08 23:03:46 +00:00
ids := []int64{}
err := pool.BeginFunc(ctx, func(tx pgx.Tx) error {
stmt := `
WITH account AS (
SELECT id
FROM accounts
2021-07-09 01:01:45 +00:00
WHERE
2021-07-09 04:27:50 +00:00
last_enqueued_at < $1
OR last_checked_at < $2
ORDER BY last_checked_at
LIMIT 1000
2021-07-08 23:03:46 +00:00
)
UPDATE accounts
2021-07-09 04:27:50 +00:00
SET last_enqueued_at = $3
2021-07-08 23:03:46 +00:00
WHERE accounts.id IN(SELECT id FROM account)
RETURNING accounts.id`
2021-07-09 04:27:50 +00:00
rows, err := tx.Query(ctx, stmt, ready, expired, now)
2021-07-08 23:03:46 +00:00
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var id int64
rows.Scan(&id)
ids = append(ids, id)
}
return nil
})
if err != nil {
logger.WithFields(logrus.Fields{
"err": err,
}).Error("failed to fetch batch of accounts")
return
}
logger.WithFields(logrus.Fields{
"count": len(ids),
2021-07-09 04:27:50 +00:00
"start": ready,
2021-07-08 23:03:46 +00:00
}).Debug("enqueueing account batch")
2021-07-09 00:17:20 +00:00
enqueued := 0
2021-07-09 00:26:01 +00:00
skipped := 0
2021-07-09 03:12:50 +00:00
// Split ids in batches
for i := 0; i < len(ids); i += batchSize {
j := i + batchSize
if j > len(ids) {
j = len(ids)
2021-07-08 23:03:46 +00:00
}
2021-07-09 03:12:50 +00:00
batch := Int64Slice(ids[i:j])
logger.WithFields(logrus.Fields{
"len": len(batch),
}).Debug("enqueueing batch")
2021-07-09 04:27:50 +00:00
lua := fmt.Sprintf(`
2021-07-09 03:12:50 +00:00
local retv={}
local ids=cjson.decode(ARGV[1])
for i=1, #ids do
local key = "locks:accounts:" .. ids[i]
if redis.call("exists", key) == 0 then
2021-07-09 04:27:50 +00:00
redis.call("setex", key, %d, 1)
2021-07-09 03:12:50 +00:00
retv[#retv + 1] = ids[i]
end
end
2021-07-09 00:26:01 +00:00
2021-07-09 03:12:50 +00:00
return retv
2021-07-09 04:27:50 +00:00
`, checkTimeout)
2021-07-09 03:12:50 +00:00
res, err := redisConn.Eval(ctx, lua, []string{}, batch).Result()
if err != nil {
2021-07-09 00:26:01 +00:00
logger.WithFields(logrus.Fields{
2021-07-09 03:12:50 +00:00
"err": err,
}).Error("failed to check for locked accounts")
2021-07-09 00:26:01 +00:00
}
2021-07-09 03:12:50 +00:00
vals := res.([]interface{})
skipped += len(batch) - len(vals)
enqueued += len(vals)
2021-07-09 03:12:50 +00:00
if len(vals) == 0 {
continue
}
batchIds := make([]string, len(vals))
for k, v := range vals {
batchIds[k] = strconv.FormatInt(v.(int64), 10)
}
if err = queue.Publish(batchIds...); err != nil {
logger.WithFields(logrus.Fields{
"err": err,
}).Error("failed to enqueue account")
}
2021-07-08 23:03:46 +00:00
}
2021-07-09 02:09:14 +00:00
statsd.Histogram("apollo.queue.enqueued", float64(enqueued), []string{}, 1)
statsd.Histogram("apollo.queue.skipped", float64(skipped), []string{}, 1)
2021-07-09 02:15:28 +00:00
statsd.Histogram("apollo.queue.runtime", float64(time.Now().Sub(start).Milliseconds()), []string{}, 1)
2021-07-09 02:09:14 +00:00
2021-07-08 23:03:46 +00:00
logger.WithFields(logrus.Fields{
2021-07-09 00:26:01 +00:00
"count": enqueued,
"skipped": skipped,
2021-07-09 04:27:50 +00:00
"start": ready,
2021-07-09 06:00:57 +00:00
}).Debug("done enqueueing account batch")
2021-07-08 23:03:46 +00:00
}
func logErrors(errChan <-chan error) {
for err := range errChan {
log.Print("error: ", err)
}
}
2021-07-09 03:12:50 +00:00
type Int64Slice []int64
func (ii Int64Slice) MarshalBinary() (data []byte, err error) {
bytes, err := json.Marshal(ii)
return bytes, err
}