SQS-service/app/billing/billing.go

// app/billing/billing.go
// Модуль учёта использования SQS-операций для биллинга.
// Записывает каждую успешную SQS-операцию в PostgreSQL: tenant_id, operation, msg_count, msg_bytes.
// Если PostgreSQL не сконфигурирован — billing отключён, SQS работает как раньше.
// Created: 2026-04-12
package billing

import (
	"database/sql"
	"fmt"
	"os"
	"sync"

	// PostgreSQL драйвер — регистрируется в database/sql через init()
	_ "github.com/lib/pq"
	log "github.com/sirupsen/logrus"
)

// db — подключение к PostgreSQL для записи usage-данных.
// nil если billing отключён. Доступ — только через getDB/setDB (мьютекс):
// usageWorker читает db параллельно с Close/Init — иначе data race и nil-паника.
var (
	db   *sql.DB
	dbMu sync.Mutex
)

// getDB — потокобезопасное чтение db.
func getDB() *sql.DB {
	dbMu.Lock()
	defer dbMu.Unlock()
	return db
}

// setDB — потокобезопасная установка db.
func setDB(d *sql.DB) {
	dbMu.Lock()
	db = d
	dbMu.Unlock()
}

// Init — подключается к PostgreSQL и создаёт таблицу sqs_usage_records если не существует.
// Env переменные: BILLING_PG_HOST, BILLING_PG_PORT, BILLING_PG_DATABASE, BILLING_PG_USER,
// BILLING_PG_PASSWORD, BILLING_PG_SSLMODE.
// Если BILLING_PG_HOST не задан — billing отключён, сервис работает без него.
func Init() {
	host := os.Getenv("BILLING_PG_HOST")
	if host == "" {
		log.Info("billing: BILLING_PG_HOST not set, usage tracking disabled")
		return
	}

	port := os.Getenv("BILLING_PG_PORT")
	if port == "" {
		port = "5432"
	}
	dbname := os.Getenv("BILLING_PG_DATABASE")
	user := os.Getenv("BILLING_PG_USER")
	password := os.Getenv("BILLING_PG_PASSWORD")
	sslmode := os.Getenv("BILLING_PG_SSLMODE")
	if sslmode == "" {
		sslmode = "require"
	}

	dsn := fmt.Sprintf("host=%s port=%s dbname=%s user=%s password=%s sslmode=%s",
		host, port, dbname, user, password, sslmode)

	var err error
	d, err := sql.Open("postgres", dsn)
	if err != nil {
		log.Errorf("billing: failed to open PostgreSQL: %v", err)
		return
	}

	// Проверяем реальное подключение (Open не подключается)
	if err = d.Ping(); err != nil {
		log.Errorf("billing: failed to connect to PostgreSQL: %v", err)
		d.Close()
		return
	}

	// Ограничиваем пул — billing не должен отжирать коннекты у основной БД
	d.SetMaxOpenConns(5)
	d.SetMaxIdleConns(2)

	setDB(d)
	if err = autoMigrate(); err != nil {
		log.Errorf("billing: failed to create table: %v", err)
		setDB(nil)
		d.Close()
		return
	}

	log.Infof("billing: connected to PostgreSQL %s:%s/%s, usage tracking enabled", host, port, dbname)
}

// autoMigrate — создаёт таблицу и индекс если не существуют.
// Идемпотентно — безопасно вызывать при каждом старте.
func autoMigrate() error {
	_, err := getDB().Exec(`
		CREATE TABLE IF NOT EXISTS sqs_usage_records (
			id          BIGSERIAL PRIMARY KEY,
			tenant_id   TEXT NOT NULL,
			operation   TEXT NOT NULL,
			queue_name  TEXT NOT NULL DEFAULT '',
			msg_count   INTEGER DEFAULT 1,
			msg_bytes   BIGINT DEFAULT 0,
			recorded_at TIMESTAMPTZ DEFAULT NOW()
		);
		CREATE INDEX IF NOT EXISTS idx_sqs_usage_tenant_time
			ON sqs_usage_records(tenant_id, recorded_at);
		-- Миграция: добавляем queue_name если таблица уже существует без него
		DO $$ BEGIN
			ALTER TABLE sqs_usage_records ADD COLUMN queue_name TEXT NOT NULL DEFAULT '';
		EXCEPTION WHEN duplicate_column THEN NULL;
		END $$;
	`)
	return err
}

// usageRecord — одна запись биллинга, накапливается в буфере воркера.
type usageRecord struct {
	tenantID  string
	operation string
	queueName string
	msgCount  int
	msgBytes  int64
}

// usageChan — канал-буфер для асинхронных записей биллинга.
// Вместо горутины на каждую операцию (навал горутин под нагрузкой) используется
// один воркер, читающий из буферизованного канала.
var usageChan = make(chan usageRecord, 1024)

// done — сигнал остановки воркера при graceful shutdown.
var done = make(chan struct{})

// closeOnce гарантирует однократный close(done).
var closeOnce sync.Once

// workerWg — ожидание завершения воркера перед закрытием db (иначе Close обнулил
// бы db, пока воркер пишет — nil-паника + потеря остатка буфера usageChan).
var workerWg sync.WaitGroup

// init запускает воркер-горутину биллинга. Воркер стартует всегда; если db == nil
// (billing отключён) — он просто дропает записи.
func init() {
	workerWg.Add(1)
	go usageWorker()
}

// usageWorker — единственная горутина, пишущая usage-записи в PostgreSQL.
func usageWorker() {
	defer workerWg.Done()
	for {
		select {
		case rec := <-usageChan:
			processUsage(rec)
		case <-done:
			// Дренируем остаток буфера перед выходом (не теряем usage).
			for {
				select {
				case rec := <-usageChan:
					processUsage(rec)
				default:
					return
				}
			}
		}
	}
}

// processUsage — одна запись usage в PostgreSQL (db == nil → дроп).
func processUsage(rec usageRecord) {
	d := getDB()
	if d == nil {
		return
	}
	_, err := d.Exec(
		`INSERT INTO sqs_usage_records (tenant_id, operation, queue_name, msg_count, msg_bytes) VALUES ($1, $2, $3, $4, $5)`,
		rec.tenantID, rec.operation, rec.queueName, rec.msgCount, rec.msgBytes,
	)
	if err != nil {
		log.Errorf("billing: failed to record usage [%s/%s]: %v", rec.tenantID, rec.operation, err)
	}
}

// RecordUsage — ставит запись в очередь биллинга (асинхронно, без блокировки SQS).
// Если billing отключён — no-op. Если буфер переполнен — запись дропается (биллинг
// жертвуем ради стабильности основного пути), ошибка логируется.
func RecordUsage(tenantID, operation, queueName string, msgCount int, msgBytes int64) {
	if getDB() == nil {
		return
	}
	select {
	case usageChan <- usageRecord{tenantID, operation, queueName, msgCount, msgBytes}:
	case <-done:
		// Воркер остановлен (shutdown) — запись не примется.
	default:
		log.Warnf("billing: буфер полон, запись [%s/%s] дропнута", tenantID, operation)
	}
}

// Enabled — возвращает true если billing подключён к PostgreSQL
func Enabled() bool {
	return getDB() != nil
}

// Close — останавливает воркер, дожидается его завершения и закрывает db.
// Вызывается при graceful shutdown. Порядок обязателен: сначала воркер,
// потом db.Close — иначе гонка db=nil с Exec воркера (nil-паника).
func Close() {
	closeOnce.Do(func() { close(done) })
	workerWg.Wait()
	dbMu.Lock()
	if db != nil {
		db.Close()
		db = nil
	}
	dbMu.Unlock()
	log.Info("billing: PostgreSQL connection closed")
}