smthqueue/cmd/simple/main.go

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.GetDBStats()
		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
}