billing: потокобезопасный db + корректный Close (остановка воркера до db.Close)

app/billing/billing.go

@@ -9,6 +9,7 @@ import (
 	"database/sql"
 	"fmt"
 	"os"
+	"sync"
 
 	// PostgreSQL драйвер — регистрируется в database/sql через init()
 	_ "github.com/lib/pq"
@@ -16,8 +17,26 @@ import (
 )
 
 // db — подключение к PostgreSQL для записи usage-данных.
-// nil если billing отключён.
-var db *sql.DB
+// 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,
@@ -46,28 +65,28 @@ func Init() {
 		host, port, dbname, user, password, sslmode)
 
 	var err error
-	db, err = sql.Open("postgres", dsn)
+	d, err := sql.Open("postgres", dsn)
 	if err != nil {
 		log.Errorf("billing: failed to open PostgreSQL: %v", err)
 		return
 	}
 
 	// Проверяем реальное подключение (Open не подключается)
-	if err = db.Ping(); err != nil {
+	if err = d.Ping(); err != nil {
 		log.Errorf("billing: failed to connect to PostgreSQL: %v", err)
-		db.Close()
-		db = nil
+		d.Close()
 		return
 	}
 
 	// Ограничиваем пул — billing не должен отжирать коннекты у основной БД
-	db.SetMaxOpenConns(5)
-	db.SetMaxIdleConns(2)
+	d.SetMaxOpenConns(5)
+	d.SetMaxIdleConns(2)
 
+	setDB(d)
 	if err = autoMigrate(); err != nil {
 		log.Errorf("billing: failed to create table: %v", err)
-		db.Close()
-		db = nil
+		setDB(nil)
+		d.Close()
 		return
 	}
 
@@ -77,7 +96,7 @@ func Init() {
 // autoMigrate — создаёт таблицу и индекс если не существуют.
 // Идемпотентно — безопасно вызывать при каждом старте.
 func autoMigrate() error {
-	_, err := db.Exec(`
+	_, err := getDB().Exec(`
 		CREATE TABLE IF NOT EXISTS sqs_usage_records (
 			id          BIGSERIAL PRIMARY KEY,
 			tenant_id   TEXT NOT NULL,
@@ -112,37 +131,70 @@ type usageRecord struct {
 // один воркер, читающий из буферизованного канала.
 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() {
-	for rec := range usageChan {
-		if db == nil {
-			continue
-		}
-		_, err := db.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)
+	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 db == nil {
+	if getDB() == nil {
 		return
 	}
 	select {
 	case usageChan <- usageRecord{tenantID, operation, queueName, msgCount, msgBytes}:
+	case <-done:
+		// Воркер остановлен (shutdown) — запись не примется.
 	default:
 		log.Warnf("billing: буфер полон, запись [%s/%s] дропнута", tenantID, operation)
 	}
@@ -150,14 +202,20 @@ func RecordUsage(tenantID, operation, queueName string, msgCount int, msgBytes i
 
 // Enabled — возвращает true если billing подключён к PostgreSQL
 func Enabled() bool {
-	return db != nil
+	return getDB() != nil
 }
 
-// Close — закрывает подключение к PostgreSQL. Вызывается при graceful shutdown.
+// 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
-		log.Info("billing: PostgreSQL connection closed")
 	}
+	dbMu.Unlock()
+	log.Info("billing: PostgreSQL connection closed")
 }