package main import ( "context" "errors" "fmt" "os" "os/signal" "runtime" "sync/atomic" "time" "git.loyso.art/frx/smthqueue/internal/postgres" "github.com/spf13/cobra" "golang.org/x/exp/slog" "golang.org/x/sync/errgroup" ) func main() { ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt) defer cancel() err := app(ctx) if err != nil { println("error running app: " + err.Error()) os.Exit(1) } } func app(ctx context.Context) error { log := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{ Level: slog.LevelInfo, })) cfg := parseEnvConfig() setupCommands(cfg, log) err := rootCommand.ExecuteContext(ctx) if err != nil { return err } return nil } func setupCommands(cfg config, log *slog.Logger) { rootCommand.AddCommand( clientCommand, ) clientCommand.AddCommand( clientConsumerCommand, clientProducerCommand, ) app := application{ cfg: cfg, log: log.With(slog.String("app", "cli")), } clientConsumerCommand.RunE = app.handleConsumerCommand clientProducerCommand.RunE = app.handleProducerCommand } var rootCommand = &cobra.Command{ Use: "smthqueue", Short: "entry point for commands", } var clientCommand = &cobra.Command{ Use: "client", Short: "a simple client for testing", } var clientConsumerCommand = &cobra.Command{ Use: "consumer [topic]", Short: "runs application in consumer mode", RunE: nil, } var clientProducerCommand = &cobra.Command{ Use: "producer [topic]", Short: "runs application in producer mode", RunE: nil, } type application struct { cfg config log *slog.Logger } type rpsReporter struct { log *slog.Logger accumulator int32 db postgres.Client } func (r *rpsReporter) inc() { atomic.AddInt32(&r.accumulator, 1) } func (r *rpsReporter) reset() int32 { return atomic.SwapInt32(&r.accumulator, 0) } func (r *rpsReporter) loop() { ticker := time.NewTicker(time.Second) for range ticker.C { rps := r.reset() dbStats := r.db.GetStats() r.log.Info( "tick", slog.Int64("rps", int64(rps)), slog.Int64("db_total_conns", int64(dbStats.TotalConnections)), slog.Int64("db_acquried_conns", int64(dbStats.AcquiredConnections)), slog.Int64("db_idle_conns", int64(dbStats.IdleConnections)), ) } } func (a *application) handleConsumerCommand(cmd *cobra.Command, args []string) error { ctx := cmd.Context() pg, err := a.getDB(ctx) if err != nil { return err } reporter := &rpsReporter{ log: a.log, db: pg, } go reporter.loop() eg, egctx := errgroup.WithContext(ctx) queueClient := pg.Queue() workers := runtime.NumCPU() a.log.InfoContext(egctx, "running workers", slog.Int("count", workers)) for i := 0; i < workers; i++ { eg.Go(func() error { for { message, err := queueClient.Dequeue(egctx, postgres.DequeueParams{ Topic: a.cfg.Topic, Timeout: time.Second, }) if err != nil { if errors.Is(err, postgres.ErrNoMessage) { time.Sleep(time.Second / 1) continue } a.log.Error("unable to dequeue message", slog.Any("err", err)) return nil } a.log.DebugContext(egctx, "scheduling new message", slog.String("message_id", message.ID)) a.log.DebugContext(egctx, "handling new message", slog.Any("message", message)) err = queueClient.Ack(egctx, message) if err != nil { a.log.ErrorContext(egctx, "unable to ack message", slog.Any("err", err)) continue } reporter.inc() a.log.DebugContext(egctx, "message acked", slog.String("message_id", message.ID)) } }) } return eg.Wait() } func (a *application) handleProducerCommand(cmd *cobra.Command, args []string) error { ctx := cmd.Context() pg, err := a.getDB(ctx) if err != nil { return err } queueClient := pg.Queue() eg, egctx := errgroup.WithContext(ctx) workers := runtime.NumCPU() for i := 0; i < workers; i++ { eg.Go(func() error { for { select { case <-egctx.Done(): return nil default: qm, err := queueClient.Enqueue(egctx, postgres.EnqueueParams{ Topic: a.cfg.Topic, Payload: []byte("simple message"), VisibleTimeout: time.Second, }) if err != nil { a.log.ErrorContext(egctx, "unable to enqueue message", slog.Any("err", err)) time.Sleep(time.Second) continue } a.log.DebugContext(egctx, "message queued", slog.Any("message", qm)) } } }) } return eg.Wait() } func (a *application) getDB(ctx context.Context) (postgres.Client, error) { pgcfg := postgres.Config{ MaxConns: 0, MaxIdleConns: 0, MasterDSN: a.cfg.DSN, } client, err := postgres.New(ctx, pgcfg, a.log) if err != nil { return nil, fmt.Errorf("making postgres client: %w", err) } return client, nil }