diff --git a/app/billing/billing.go b/app/billing/billing.go index caab50c..ef946b3 100644 --- a/app/billing/billing.go +++ b/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") }