| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394 |
- package repository
- import (
- "context"
- "database/sql"
- "encoding/json"
- "errors"
- "fmt"
- "strconv"
- "strings"
- "sync/atomic"
- "time"
- dbent "github.com/Wei-Shaw/sub2api/ent"
- "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
- "github.com/Wei-Shaw/sub2api/internal/service"
- )
- // usageLogInsertArgTypes must stay in the same order as:
- // 1. prepareUsageLogInsert().args
- // 2. every INSERT/CTE VALUES column list in this file
- // 3. execUsageLogInsertNoResult placeholder positions
- // 4. scanUsageLog selected column order (via usageLogSelectColumns)
- //
- // When adding a usage_logs column, update all of those call sites together.
- var usageLogInsertArgTypes = [...]string{
- "bigint", // user_id
- "bigint", // api_key_id
- "bigint", // account_id
- "text", // request_id
- "text", // model
- "text", // requested_model
- "text", // upstream_model
- "text", // upstream_response_model
- "boolean", // upstream_model_mismatch
- "bigint", // group_id
- "bigint", // subscription_id
- "integer", // input_tokens
- "integer", // output_tokens
- "integer", // cache_creation_tokens
- "integer", // cache_read_tokens
- "integer", // cache_creation_5m_tokens
- "integer", // cache_creation_1h_tokens
- "integer", // image_output_tokens
- "numeric", // image_output_cost
- "integer", // image_input_tokens
- "numeric", // image_input_cost
- "numeric", // input_cost
- "numeric", // output_cost
- "numeric", // cache_creation_cost
- "numeric", // cache_read_cost
- "numeric", // total_cost
- "numeric", // actual_cost
- "numeric", // rate_multiplier
- "numeric", // account_rate_multiplier
- "smallint", // billing_type
- "smallint", // request_type
- "boolean", // stream
- "boolean", // openai_ws_mode
- "integer", // duration_ms
- "integer", // first_token_ms
- "text", // user_agent
- "text", // ip_address
- "integer", // image_count
- "text", // image_size
- "text", // image_input_size
- "text", // image_output_size
- "text", // image_size_source
- "jsonb", // image_size_breakdown
- "integer", // video_count
- "text", // video_resolution
- "integer", // video_duration_seconds
- "text", // service_tier
- "text", // reasoning_effort
- "text", // requested_reasoning_effort
- "text", // inbound_endpoint
- "text", // upstream_endpoint
- "boolean", // cache_ttl_overridden
- "boolean", // long_context_billing_applied
- "bigint", // channel_id
- "text", // model_mapping_chain
- "text", // billing_tier
- "text", // billing_mode
- "numeric", // account_stats_cost
- "text", // upstream_request_id
- "text", // session_id
- "boolean", // native_compaction_v2
- "timestamptz", // created_at
- }
- const (
- usageLogCreateBatchMaxSize = 64
- usageLogCreateBatchWindow = 3 * time.Millisecond
- usageLogCreateBatchQueueCap = 4096
- usageLogCreateCancelWait = 2 * time.Second
- usageLogBestEffortBatchMaxSize = 256
- usageLogBestEffortBatchWindow = 20 * time.Millisecond
- usageLogBestEffortBatchQueueCap = 32768
- usageLogBestEffortRecentTTL = 30 * time.Second
- )
- type usageLogCreateRequest struct {
- log *service.UsageLog
- prepared usageLogInsertPrepared
- shared *usageLogCreateShared
- resultCh chan usageLogCreateResult
- }
- type usageLogCreateResult struct {
- inserted bool
- err error
- }
- type usageLogBestEffortRequest struct {
- prepared usageLogInsertPrepared
- apiKeyID int64
- resultCh chan error
- }
- type usageLogInsertPrepared struct {
- createdAt time.Time
- requestID string
- rateMultiplier float64
- requestType int16
- args []any
- }
- type usageLogBatchState struct {
- ID int64
- CreatedAt time.Time
- }
- type usageLogBatchRow struct {
- RequestID string `json:"request_id"`
- APIKeyID int64 `json:"api_key_id"`
- ID int64 `json:"id"`
- CreatedAt time.Time `json:"created_at"`
- Inserted bool `json:"inserted"`
- }
- type usageLogCreateShared struct {
- state atomic.Int32
- }
- const (
- usageLogCreateStateQueued int32 = iota
- usageLogCreateStateProcessing
- usageLogCreateStateCompleted
- usageLogCreateStateCanceled
- )
- func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
- if log == nil {
- return false, nil
- }
- if tx := dbent.TxFromContext(ctx); tx != nil {
- return r.createSingle(ctx, tx.Client(), log)
- }
- requestID := strings.TrimSpace(log.RequestID)
- if requestID == "" {
- return r.createSingle(ctx, r.sql, log)
- }
- log.RequestID = requestID
- return r.createBatched(ctx, log)
- }
- func (r *usageLogRepository) CreateBestEffort(ctx context.Context, log *service.UsageLog) error {
- if log == nil {
- return nil
- }
- if tx := dbent.TxFromContext(ctx); tx != nil {
- _, err := r.createSingle(ctx, tx.Client(), log)
- return err
- }
- if r.db == nil {
- _, err := r.createSingle(ctx, r.sql, log)
- return err
- }
- r.ensureBestEffortBatcher()
- if r.bestEffortBatchCh == nil {
- _, err := r.createSingle(ctx, r.sql, log)
- return err
- }
- req := usageLogBestEffortRequest{
- prepared: prepareUsageLogInsert(log),
- apiKeyID: log.APIKeyID,
- resultCh: make(chan error, 1),
- }
- if key, ok := r.bestEffortRecentKey(req.prepared.requestID, req.apiKeyID); ok {
- if _, exists := r.bestEffortRecent.Get(key); exists {
- return nil
- }
- }
- // 队列满时阻塞等待而非立即丢弃:批处理器持续排空队列,短暂等待即可入队。
- // 立即丢弃会造成“已扣费但无 usage_log”的永久数据缺口(issue #3656);
- // 阻塞上限由调用方 ctx 期限约束,超时后由上层同步兜底。
- select {
- case r.bestEffortBatchCh <- req:
- case <-ctx.Done():
- return service.MarkUsageLogCreateDropped(ctx.Err())
- }
- select {
- case err := <-req.resultCh:
- return err
- case <-ctx.Done():
- return service.MarkUsageLogCreateDropped(ctx.Err())
- }
- }
- func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor, log *service.UsageLog) (bool, error) {
- prepared := prepareUsageLogInsert(log)
- if sqlq == nil {
- sqlq = r.sql
- }
- if ctx != nil && ctx.Err() != nil {
- return false, service.MarkUsageLogCreateNotPersisted(ctx.Err())
- }
- query := `
- INSERT INTO usage_logs (
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- ) VALUES (
- $1, $2, $3, $4, $5, $6, $7, $8, $9,
- $10, $11,
- $12, $13, $14, $15,
- $16, $17, $18, $19,
- $20, $21, $22, $23, $24, $25,
- $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59, $60, $61, $62
- )
- ON CONFLICT (request_id, api_key_id) DO NOTHING
- RETURNING id, created_at
- `
- if err := scanSingleRow(ctx, sqlq, query, prepared.args, &log.ID, &log.CreatedAt); err != nil {
- if errors.Is(err, sql.ErrNoRows) && prepared.requestID != "" {
- selectQuery := "SELECT id, created_at FROM usage_logs WHERE request_id = $1 AND api_key_id = $2"
- if err := scanSingleRow(ctx, sqlq, selectQuery, []any{prepared.requestID, log.APIKeyID}, &log.ID, &log.CreatedAt); err != nil {
- return false, err
- }
- log.RateMultiplier = prepared.rateMultiplier
- return false, nil
- } else {
- return false, err
- }
- }
- log.RateMultiplier = prepared.rateMultiplier
- return true, nil
- }
- func (r *usageLogRepository) createBatched(ctx context.Context, log *service.UsageLog) (bool, error) {
- if r.db == nil {
- return r.createSingle(ctx, r.sql, log)
- }
- r.ensureCreateBatcher()
- if r.createBatchCh == nil {
- return r.createSingle(ctx, r.sql, log)
- }
- req := usageLogCreateRequest{
- log: log,
- prepared: prepareUsageLogInsert(log),
- shared: &usageLogCreateShared{},
- resultCh: make(chan usageLogCreateResult, 1),
- }
- // 队列满时阻塞等待而非立即报错:本路径是 best-effort 丢弃后的最后兜底,
- // 立即失败会让日志永久丢失;阻塞上限由调用方 ctx 期限约束。
- select {
- case r.createBatchCh <- req:
- case <-ctx.Done():
- return false, service.MarkUsageLogCreateNotPersisted(ctx.Err())
- }
- select {
- case res := <-req.resultCh:
- return res.inserted, res.err
- case <-ctx.Done():
- if req.shared != nil && req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateCanceled) {
- return false, service.MarkUsageLogCreateNotPersisted(ctx.Err())
- }
- timer := time.NewTimer(usageLogCreateCancelWait)
- defer timer.Stop()
- select {
- case res := <-req.resultCh:
- return res.inserted, res.err
- case <-timer.C:
- return false, ctx.Err()
- }
- }
- }
- func (r *usageLogRepository) ensureCreateBatcher() {
- if r == nil || r.db == nil {
- return
- }
- // nil 检查必须在 Once 内部:在外层做无同步快路径读会与 Once 内的写构成数据竞争。
- r.createBatchOnce.Do(func() {
- if r.createBatchCh == nil {
- r.createBatchCh = make(chan usageLogCreateRequest, usageLogCreateBatchQueueCap)
- go r.runCreateBatcher(r.db)
- }
- })
- }
- func (r *usageLogRepository) ensureBestEffortBatcher() {
- if r == nil || r.db == nil {
- return
- }
- // 同 ensureCreateBatcher:nil 检查放在 Once 内部以避免数据竞争。
- r.bestEffortBatchOnce.Do(func() {
- if r.bestEffortBatchCh == nil {
- r.bestEffortBatchCh = make(chan usageLogBestEffortRequest, usageLogBestEffortBatchQueueCap)
- go r.runBestEffortBatcher(r.db)
- }
- })
- }
- func (r *usageLogRepository) runCreateBatcher(db *sql.DB) {
- for {
- first, ok := <-r.createBatchCh
- if !ok {
- return
- }
- batch := make([]usageLogCreateRequest, 0, usageLogCreateBatchMaxSize)
- batch = append(batch, first)
- timer := time.NewTimer(usageLogCreateBatchWindow)
- batchLoop:
- for len(batch) < usageLogCreateBatchMaxSize {
- select {
- case req, ok := <-r.createBatchCh:
- if !ok {
- break batchLoop
- }
- batch = append(batch, req)
- case <-timer.C:
- break batchLoop
- }
- }
- if !timer.Stop() {
- select {
- case <-timer.C:
- default:
- }
- }
- r.flushCreateBatch(db, batch)
- }
- }
- func (r *usageLogRepository) runBestEffortBatcher(db *sql.DB) {
- for {
- first, ok := <-r.bestEffortBatchCh
- if !ok {
- return
- }
- batch := make([]usageLogBestEffortRequest, 0, usageLogBestEffortBatchMaxSize)
- batch = append(batch, first)
- timer := time.NewTimer(usageLogBestEffortBatchWindow)
- bestEffortLoop:
- for len(batch) < usageLogBestEffortBatchMaxSize {
- select {
- case req, ok := <-r.bestEffortBatchCh:
- if !ok {
- break bestEffortLoop
- }
- batch = append(batch, req)
- case <-timer.C:
- break bestEffortLoop
- }
- }
- if !timer.Stop() {
- select {
- case <-timer.C:
- default:
- }
- }
- r.flushBestEffortBatch(db, batch)
- }
- }
- func (r *usageLogRepository) flushCreateBatch(db *sql.DB, batch []usageLogCreateRequest) {
- if len(batch) == 0 {
- return
- }
- uniqueOrder := make([]string, 0, len(batch))
- preparedByKey := make(map[string]usageLogInsertPrepared, len(batch))
- requestsByKey := make(map[string][]usageLogCreateRequest, len(batch))
- fallback := make([]usageLogCreateRequest, 0)
- for _, req := range batch {
- if req.log == nil {
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
- continue
- }
- if req.shared != nil && !req.shared.state.CompareAndSwap(usageLogCreateStateQueued, usageLogCreateStateProcessing) {
- if req.shared.state.Load() == usageLogCreateStateCanceled {
- completeUsageLogCreateRequest(req, usageLogCreateResult{
- inserted: false,
- err: service.MarkUsageLogCreateNotPersisted(context.Canceled),
- })
- continue
- }
- }
- prepared := req.prepared
- if prepared.requestID == "" {
- fallback = append(fallback, req)
- continue
- }
- key := usageLogBatchKey(prepared.requestID, req.log.APIKeyID)
- if _, exists := requestsByKey[key]; !exists {
- uniqueOrder = append(uniqueOrder, key)
- preparedByKey[key] = prepared
- }
- requestsByKey[key] = append(requestsByKey[key], req)
- }
- if len(uniqueOrder) > 0 {
- insertedMap, stateMap, safeFallback, err := r.batchInsertUsageLogs(db, uniqueOrder, preparedByKey)
- if err != nil {
- if safeFallback {
- for _, key := range uniqueOrder {
- fallback = append(fallback, requestsByKey[key]...)
- }
- } else {
- for _, key := range uniqueOrder {
- reqs := requestsByKey[key]
- state, hasState := stateMap[key]
- inserted := insertedMap[key]
- for idx, req := range reqs {
- req.log.RateMultiplier = preparedByKey[key].rateMultiplier
- if hasState {
- req.log.ID = state.ID
- req.log.CreatedAt = state.CreatedAt
- }
- switch {
- case inserted && idx == 0:
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: true, err: nil})
- case inserted:
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
- case hasState:
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
- case idx == 0:
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: err})
- default:
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: false, err: nil})
- }
- }
- }
- }
- } else {
- for _, key := range uniqueOrder {
- reqs := requestsByKey[key]
- state, ok := stateMap[key]
- if !ok {
- for _, req := range reqs {
- completeUsageLogCreateRequest(req, usageLogCreateResult{
- inserted: false,
- err: fmt.Errorf("usage log batch state missing for key=%s", key),
- })
- }
- continue
- }
- for idx, req := range reqs {
- req.log.ID = state.ID
- req.log.CreatedAt = state.CreatedAt
- req.log.RateMultiplier = preparedByKey[key].rateMultiplier
- completeUsageLogCreateRequest(req, usageLogCreateResult{
- inserted: idx == 0 && insertedMap[key],
- err: nil,
- })
- }
- }
- }
- }
- if len(fallback) == 0 {
- return
- }
- for _, req := range fallback {
- fallbackCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
- inserted, err := r.createSingle(fallbackCtx, db, req.log)
- cancel()
- completeUsageLogCreateRequest(req, usageLogCreateResult{inserted: inserted, err: err})
- }
- }
- func (r *usageLogRepository) flushBestEffortBatch(db *sql.DB, batch []usageLogBestEffortRequest) {
- if len(batch) == 0 {
- return
- }
- type bestEffortGroup struct {
- prepared usageLogInsertPrepared
- apiKeyID int64
- key string
- reqs []usageLogBestEffortRequest
- }
- groupsByKey := make(map[string]*bestEffortGroup, len(batch))
- groupOrder := make([]*bestEffortGroup, 0, len(batch))
- preparedList := make([]usageLogInsertPrepared, 0, len(batch))
- for idx, req := range batch {
- prepared := req.prepared
- key := fmt.Sprintf("__best_effort_%d", idx)
- if prepared.requestID != "" {
- key = usageLogBatchKey(prepared.requestID, req.apiKeyID)
- }
- group, exists := groupsByKey[key]
- if !exists {
- group = &bestEffortGroup{
- prepared: prepared,
- apiKeyID: req.apiKeyID,
- key: key,
- }
- groupsByKey[key] = group
- groupOrder = append(groupOrder, group)
- preparedList = append(preparedList, prepared)
- }
- group.reqs = append(group.reqs, req)
- }
- if len(preparedList) == 0 {
- for _, req := range batch {
- sendUsageLogBestEffortResult(req.resultCh, nil)
- }
- return
- }
- ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
- defer cancel()
- query, args := buildUsageLogBestEffortInsertQuery(preparedList)
- if _, err := db.ExecContext(ctx, query, args...); err != nil {
- logger.LegacyPrintf("repository.usage_log", "best-effort batch insert failed: %v", err)
- for _, group := range groupOrder {
- singleErr := execUsageLogInsertNoResult(ctx, db, group.prepared)
- if singleErr != nil {
- logger.LegacyPrintf("repository.usage_log", "best-effort single fallback insert failed: %v", singleErr)
- } else if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil {
- r.bestEffortRecent.SetDefault(group.key, struct{}{})
- }
- for _, req := range group.reqs {
- sendUsageLogBestEffortResult(req.resultCh, singleErr)
- }
- }
- return
- }
- for _, group := range groupOrder {
- if group.prepared.requestID != "" && r != nil && r.bestEffortRecent != nil {
- r.bestEffortRecent.SetDefault(group.key, struct{}{})
- }
- for _, req := range group.reqs {
- sendUsageLogBestEffortResult(req.resultCh, nil)
- }
- }
- }
- func sendUsageLogBestEffortResult(ch chan error, err error) {
- if ch == nil {
- return
- }
- select {
- case ch <- err:
- default:
- }
- }
- func completeUsageLogCreateRequest(req usageLogCreateRequest, res usageLogCreateResult) {
- if req.shared != nil {
- req.shared.state.Store(usageLogCreateStateCompleted)
- }
- sendUsageLogCreateResult(req.resultCh, res)
- }
- func (r *usageLogRepository) batchInsertUsageLogs(db *sql.DB, keys []string, preparedByKey map[string]usageLogInsertPrepared) (map[string]bool, map[string]usageLogBatchState, bool, error) {
- if len(keys) == 0 {
- return map[string]bool{}, map[string]usageLogBatchState{}, false, nil
- }
- ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
- defer cancel()
- query, args := buildUsageLogBatchInsertQuery(keys, preparedByKey)
- var payload []byte
- if err := db.QueryRowContext(ctx, query, args...).Scan(&payload); err != nil {
- return nil, nil, true, err
- }
- var rows []usageLogBatchRow
- if err := json.Unmarshal(payload, &rows); err != nil {
- return nil, nil, false, err
- }
- insertedMap := make(map[string]bool, len(keys))
- stateMap := make(map[string]usageLogBatchState, len(keys))
- for _, row := range rows {
- key := usageLogBatchKey(row.RequestID, row.APIKeyID)
- insertedMap[key] = row.Inserted
- stateMap[key] = usageLogBatchState{
- ID: row.ID,
- CreatedAt: row.CreatedAt,
- }
- }
- if len(stateMap) != len(keys) {
- return insertedMap, stateMap, false, fmt.Errorf("usage log batch state count mismatch: got=%d want=%d", len(stateMap), len(keys))
- }
- return insertedMap, stateMap, false, nil
- }
- func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usageLogInsertPrepared) (string, []any) {
- var query strings.Builder
- _, _ = query.WriteString(`
- WITH input (
- input_idx,
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- ) AS (VALUES `)
- // Each batch row prepends the synthetic input_index before the 60
- // usage-log column values.
- args := make([]any, 0, len(keys)*61)
- argPos := 1
- for idx, key := range keys {
- if idx > 0 {
- _, _ = query.WriteString(",")
- }
- _, _ = query.WriteString("(")
- _, _ = query.WriteString("$")
- _, _ = query.WriteString(strconv.Itoa(argPos))
- args = append(args, idx)
- argPos++
- prepared := preparedByKey[key]
- for i := 0; i < len(prepared.args); i++ {
- _, _ = query.WriteString(",")
- _, _ = query.WriteString("$")
- _, _ = query.WriteString(strconv.Itoa(argPos))
- if i < len(usageLogInsertArgTypes) {
- _, _ = query.WriteString("::")
- _, _ = query.WriteString(usageLogInsertArgTypes[i])
- }
- argPos++
- }
- _, _ = query.WriteString(")")
- args = append(args, prepared.args...)
- }
- _, _ = query.WriteString(`
- ),
- inserted AS (
- INSERT INTO usage_logs (
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- )
- SELECT
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- FROM input
- ON CONFLICT (request_id, api_key_id) DO NOTHING
- RETURNING request_id, api_key_id, id, created_at
- ),
- resolved AS (
- SELECT
- input.input_idx,
- input.request_id,
- input.api_key_id,
- COALESCE(inserted.id, existing.id) AS id,
- COALESCE(inserted.created_at, existing.created_at) AS created_at,
- (inserted.id IS NOT NULL) AS inserted
- FROM input
- LEFT JOIN inserted
- ON inserted.request_id = input.request_id
- AND inserted.api_key_id = input.api_key_id
- LEFT JOIN usage_logs existing
- ON existing.request_id = input.request_id
- AND existing.api_key_id = input.api_key_id
- )
- SELECT COALESCE(
- json_agg(
- json_build_object(
- 'request_id', resolved.request_id,
- 'api_key_id', resolved.api_key_id,
- 'id', resolved.id,
- 'created_at', resolved.created_at,
- 'inserted', resolved.inserted
- )
- ORDER BY resolved.input_idx
- ),
- '[]'::json
- )
- FROM resolved
- `)
- return query.String(), args
- }
- func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (string, []any) {
- var query strings.Builder
- _, _ = query.WriteString(`
- WITH input (
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- ) AS (VALUES `)
- args := make([]any, 0, len(preparedList)*60)
- argPos := 1
- for idx, prepared := range preparedList {
- if idx > 0 {
- _, _ = query.WriteString(",")
- }
- _, _ = query.WriteString("(")
- for i := 0; i < len(prepared.args); i++ {
- if i > 0 {
- _, _ = query.WriteString(",")
- }
- _, _ = query.WriteString("$")
- _, _ = query.WriteString(strconv.Itoa(argPos))
- if i < len(usageLogInsertArgTypes) {
- _, _ = query.WriteString("::")
- _, _ = query.WriteString(usageLogInsertArgTypes[i])
- }
- argPos++
- }
- _, _ = query.WriteString(")")
- args = append(args, prepared.args...)
- }
- _, _ = query.WriteString(`
- )
- INSERT INTO usage_logs (
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- )
- SELECT
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- FROM input
- ON CONFLICT (request_id, api_key_id) DO NOTHING
- `)
- return query.String(), args
- }
- func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared usageLogInsertPrepared) error {
- _, err := sqlq.ExecContext(ctx, `
- INSERT INTO usage_logs (
- user_id,
- api_key_id,
- account_id,
- request_id,
- model,
- requested_model,
- upstream_model,
- upstream_response_model,
- upstream_model_mismatch,
- group_id,
- subscription_id,
- input_tokens,
- output_tokens,
- cache_creation_tokens,
- cache_read_tokens,
- cache_creation_5m_tokens,
- cache_creation_1h_tokens,
- image_output_tokens,
- image_output_cost,
- image_input_tokens,
- image_input_cost,
- input_cost,
- output_cost,
- cache_creation_cost,
- cache_read_cost,
- total_cost,
- actual_cost,
- rate_multiplier,
- account_rate_multiplier,
- billing_type,
- request_type,
- stream,
- openai_ws_mode,
- duration_ms,
- first_token_ms,
- user_agent,
- ip_address,
- image_count,
- image_size,
- image_input_size,
- image_output_size,
- image_size_source,
- image_size_breakdown,
- video_count,
- video_resolution,
- video_duration_seconds,
- service_tier,
- reasoning_effort,
- requested_reasoning_effort,
- inbound_endpoint,
- upstream_endpoint,
- cache_ttl_overridden,
- long_context_billing_applied,
- channel_id,
- model_mapping_chain,
- billing_tier,
- billing_mode,
- account_stats_cost,
- upstream_request_id,
- session_id,
- native_compaction_v2,
- created_at
- ) VALUES (
- $1, $2, $3, $4, $5, $6, $7, $8, $9,
- $10, $11,
- $12, $13, $14, $15,
- $16, $17, $18, $19,
- $20, $21, $22, $23, $24, $25,
- $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57, $58, $59, $60, $61, $62
- )
- ON CONFLICT (request_id, api_key_id) DO NOTHING
- `, prepared.args...)
- return err
- }
- func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
- createdAt := log.CreatedAt
- if createdAt.IsZero() {
- createdAt = time.Now()
- }
- requestID := strings.TrimSpace(log.RequestID)
- log.RequestID = requestID
- rateMultiplier := log.RateMultiplier
- log.SyncRequestTypeAndLegacyFields()
- requestType := int16(log.RequestType)
- groupID := nullInt64(log.GroupID)
- subscriptionID := nullInt64(log.SubscriptionID)
- duration := nullInt(log.DurationMs)
- firstToken := nullInt(log.FirstTokenMs)
- userAgent := nullString(log.UserAgent)
- ipAddress := nullString(log.IPAddress)
- imageSize := nullString(log.ImageSize)
- imageInputSize := nullString(log.ImageInputSize)
- imageOutputSize := nullString(log.ImageOutputSize)
- imageSizeSource := nullString(log.ImageSizeSource)
- imageSizeBreakdown := nullStringIntMapJSON(log.ImageSizeBreakdown)
- videoResolution := nullString(log.VideoResolution)
- videoDurationSeconds := nullInt(log.VideoDurationSeconds)
- serviceTier := nullString(log.ServiceTier)
- reasoningEffort := nullString(log.ReasoningEffort)
- requestedReasoningEffort := nullString(log.RequestedReasoningEffort)
- inboundEndpoint := nullString(log.InboundEndpoint)
- upstreamEndpoint := nullString(log.UpstreamEndpoint)
- channelID := nullInt64(log.ChannelID)
- modelMappingChain := nullString(log.ModelMappingChain)
- billingTier := nullString(log.BillingTier)
- billingMode := nullString(log.BillingMode)
- upstreamRequestID := nullString(log.UpstreamRequestID)
- sessionID := nullString(log.SessionID)
- requestedModel := strings.TrimSpace(log.RequestedModel)
- if requestedModel == "" {
- requestedModel = strings.TrimSpace(log.Model)
- }
- upstreamModel := nullString(log.UpstreamModel)
- upstreamResponseModel := nullString(log.UpstreamResponseModel)
- upstreamModelMismatch := nullBool(log.UpstreamModelMismatch)
- var requestIDArg any
- if requestID != "" {
- requestIDArg = requestID
- }
- return usageLogInsertPrepared{
- createdAt: createdAt,
- requestID: requestID,
- rateMultiplier: rateMultiplier,
- requestType: requestType,
- args: []any{
- log.UserID,
- log.APIKeyID,
- log.AccountID,
- requestIDArg,
- log.Model,
- nullString(&requestedModel),
- upstreamModel,
- upstreamResponseModel,
- upstreamModelMismatch,
- groupID,
- subscriptionID,
- log.InputTokens,
- log.OutputTokens,
- log.CacheCreationTokens,
- log.CacheReadTokens,
- log.CacheCreation5mTokens,
- log.CacheCreation1hTokens,
- log.ImageOutputTokens,
- log.ImageOutputCost,
- log.ImageInputTokens,
- log.ImageInputCost,
- log.InputCost,
- log.OutputCost,
- log.CacheCreationCost,
- log.CacheReadCost,
- log.TotalCost,
- log.ActualCost,
- rateMultiplier,
- log.AccountRateMultiplier,
- log.BillingType,
- requestType,
- log.Stream,
- log.OpenAIWSMode,
- duration,
- firstToken,
- userAgent,
- ipAddress,
- log.ImageCount,
- imageSize,
- imageInputSize,
- imageOutputSize,
- imageSizeSource,
- imageSizeBreakdown,
- log.VideoCount,
- videoResolution,
- videoDurationSeconds,
- serviceTier,
- reasoningEffort,
- requestedReasoningEffort,
- inboundEndpoint,
- upstreamEndpoint,
- log.CacheTTLOverridden,
- log.LongContextBillingApplied,
- channelID,
- modelMappingChain,
- billingTier,
- billingMode,
- log.AccountStatsCost, // account_stats_cost
- upstreamRequestID, // upstream_request_id
- sessionID, // session_id
- log.NativeCompactionV2,
- createdAt,
- },
- }
- }
- func usageLogBatchKey(requestID string, apiKeyID int64) string {
- return requestID + "\x1f" + strconv.FormatInt(apiKeyID, 10)
- }
- func sendUsageLogCreateResult(ch chan usageLogCreateResult, res usageLogCreateResult) {
- if ch == nil {
- return
- }
- select {
- case ch <- res:
- default:
- }
- }
- func (r *usageLogRepository) bestEffortRecentKey(requestID string, apiKeyID int64) (string, bool) {
- requestID = strings.TrimSpace(requestID)
- if requestID == "" || r == nil || r.bestEffortRecent == nil {
- return "", false
- }
- return usageLogBatchKey(requestID, apiKeyID), true
- }
|