fix async writes, eviction

1 files changed, 37 insertions, 6 deletions

diff --git a/internal/cache/cache.go b/internal/cache/cache.go
index d6a31f3..8b13dd1 100644
--- a/internal/cache/cache.go
+++ b/internal/cache/cache.go
@@ -40,12 +40,17 @@ type Cache struct {
 	entries map[Key]*entry
 	maxSize int
 	db *sql.DB
-
+	dbCh chan dbWrite
+	wg sync.WaitGroup
 	hits int64
 	misses int64
 	evicted int64
 	stopCh chan struct{}
 }
+type dbWrite struct {
+	key Key
+	e *entry
+}
 
 func NewCache(maxSize int, dbPath string) (*Cache, error) {
 	if maxSize <= 0 {
@@ -84,6 +89,9 @@ func NewCache(maxSize int, dbPath string) (*Cache, error) {
 		}
 
 		c.loadFromDB()
+		c.dbCh = make(chan dbWrite, 1024)
+		c.wg.Add(1)
+		go c.dbWriter()
 	}
 	go c.evictLoop()
 	return c, nil
@@ -92,6 +100,10 @@ func NewCache(maxSize int, dbPath string) (*Cache, error) {
 func (c *Cache) Stop() {
 	close(c.stopCh)
 	if c.db != nil {
+		close(c.dbCh)
+	}
+	c.wg.Wait()
+	if c.db != nil {
 		c.db.Close()
 	}
 }
@@ -135,7 +147,10 @@ func (c *Cache) Set(key Key, msg *dns.Msg, ttl time.Duration) {
 	c.mu.Unlock()
 
 	if c.db != nil {
-		c.writeToDB(key, e)
+		select {
+		case c.dbCh <- dbWrite{key: key, e: e}:
+		default:
+		}
 	}
 }
 
@@ -185,15 +200,25 @@ func (c *Cache) evictLoop() {
 	for {
 		select {
 		case <-tk.C:
-			c.mu.Lock()
 			now := time.Now()
+			c.mu.RLock()
+			var expired []Key
 			for k, e := range c.entries {
 				if now.After(e.storedAt.Add(e.ttl)) {
-					delete(c.entries, k)
-					atomic.AddInt64(&c.evicted, 1)
+					expired = append(expired, k)
+				}
+			}
+			c.mu.RUnlock()
+			if len(expired) > 0 {
+				c.mu.Lock()
+				for _, k := range expired {
+					if _, ok := c.entries[k]; ok {
+						delete(c.entries, k)
+						atomic.AddInt64(&c.evicted, 1)
+					}
 				}
+				c.mu.Unlock()
 			}
-			c.mu.Unlock()
 		case <-c.stopCh:
 			return
 		}
@@ -218,6 +243,12 @@ func (c *Cache) writeToDB(key Key, e *entry) {
 	}
 }
 
+func (c *Cache) dbWriter() {
+	defer c.wg.Done()
+	for w := range c.dbCh {
+		c.writeToDB(w.key, w.e)
+	}
+}
 func (c *Cache) loadFromDB() {
 	rows, err := c.db.Query(
 		`SELECT name, qtype, class, data, stored_at, ttl_ns FROM cache`,