package imgcache import ( "context" "fmt" "sync/atomic" ) // StartPendingCountWrites starts the goroutine that writes the pending // counts to the database whenever there are some, one UPDATE at a time: // however many requests pass their deadline before their counts are // written, at most this one write waits for the database. It is a no-op // when already started. The goroutine outlives the caller, so it runs with // its own context, which StopPendingCountWrites cancels to tell it to // finish. func (c *Cache) StartPendingCountWrites() { if c.pendingCountsCancel != nil { return } ctx, cancel := context.WithCancel(context.Background()) c.pendingCountsCancel = cancel go func() { c.pendingCountWriteLoop(ctx) }() } // StopPendingCountWrites tells the goroutine StartPendingCountWrites // started, if it did, to finish, and waits for its write under way to end // and for it to return. It then writes the pending counts left, so that // they are written before the database closes. It waits for all of this at // most until ctx ends, then logs the counts not written and returns an // error. func (c *Cache) StopPendingCountWrites(ctx context.Context) error { var err error if c.pendingCountsCancel != nil { c.pendingCountsCancel() select { case <-c.pendingCountsDone: case <-ctx.Done(): // The write under way is left to finish: the database's // close waits for a query under way. err = fmt.Errorf("pending counts still being written: %w", ctx.Err()) } } if err == nil { err = c.writePendingCounts(ctx) } if err != nil { c.log.Error("counts not written at shutdown", "hits", c.pendingHits.Load(), "misses", c.pendingMisses.Load(), "upstream_fetches", c.pendingUpstreamFetches.Load(), "upstream_fetch_bytes", c.pendingUpstreamFetchBytes.Load(), "transforms", c.pendingTransforms.Load(), "error", err, ) return err } return nil } // pendingCountWriteLoop is the body of the goroutine // StartPendingCountWrites starts. It returns when ctx is cancelled, but // does not cut off a write under way then. A write that fails leaves the // counts pending, for the next write. func (c *Cache) pendingCountWriteLoop(ctx context.Context) { defer close(c.pendingCountsDone) for { select { case <-ctx.Done(): return case <-c.pendingCountsAdded: } err := c.writePendingCounts(context.WithoutCancel(ctx)) if err != nil { c.log.Warn("failed to write pending counts", "error", err) } } } // addPendingCount adds n to pendingCount, one of the pending counts, and // wakes the goroutine that writes them. pendingCountsAdded has capacity one, // so a wakeup already waiting is enough. func (c *Cache) addPendingCount(pendingCount *atomic.Int64, n int64) { pendingCount.Add(n) select { case c.pendingCountsAdded <- struct{}{}: default: } } // writePendingCounts adds the pending counts to the cache_stats row in one // UPDATE, then takes what it wrote out of them, leaving any added // meanwhile. When the UPDATE fails, they all stay pending. func (c *Cache) writePendingCounts(ctx context.Context) error { c.pendingCountsWriteMutex.Lock() defer c.pendingCountsWriteMutex.Unlock() hits := c.pendingHits.Load() misses := c.pendingMisses.Load() upstreamFetches := c.pendingUpstreamFetches.Load() upstreamFetchBytes := c.pendingUpstreamFetchBytes.Load() transforms := c.pendingTransforms.Load() if hits+misses+upstreamFetches+upstreamFetchBytes+transforms == 0 { return nil } _, err := c.db.ExecContext(ctx, ` UPDATE cache_stats SET hit_count = hit_count + ?, miss_count = miss_count + ?, upstream_fetch_count = upstream_fetch_count + ?, upstream_fetch_bytes = upstream_fetch_bytes + ?, transform_count = transform_count + ?, last_updated_at = CURRENT_TIMESTAMP WHERE id = 1 `, hits, misses, upstreamFetches, upstreamFetchBytes, transforms) if err != nil { return fmt.Errorf("failed to write pending counts: %w", err) } c.pendingHits.Add(-hits) c.pendingMisses.Add(-misses) c.pendingUpstreamFetches.Add(-upstreamFetches) c.pendingUpstreamFetchBytes.Add(-upstreamFetchBytes) c.pendingTransforms.Add(-transforms) return nil }