acme-dns-httpreq/db.go

package main

import (
	"sync"
	"time"
)

func init() {
	db = NewDB()
}

type DB struct {
	data    map[string]map[string]time.Time
	lock    *sync.Mutex
	timeout time.Duration
}

func (db *DB) Lookup(name string) []string {
	now := time.Now()
	db.lock.Lock()
	defer db.lock.Unlock()

	if domain, ok := db.data[name]; ok {
		if len(domain) > 0 {
			rv := make([]string, 0, len(domain))
			for value, expiresAt := range domain {
				if expiresAt.After(now) {
					rv = append(rv, value)
				} else {
					delete(domain, value)
				}
			}
			return rv
		} else {
			delete(db.data, name)
		}
	}
	return []string{}
}

func (db *DB) Add(name, value string) {
	expiresAt := time.Now().Add(db.timeout)
	db.lock.Lock()
	defer db.lock.Unlock()

	if _, ok := db.data[name]; ok {
		db.data[name][value] = expiresAt
	} else {
		db.data[name] = map[string]time.Time{value: expiresAt}
	}
}

func (db *DB) Remove(name, value string) {
	db.lock.Lock()
	defer db.lock.Unlock()

	if domain, ok := db.data[name]; ok {
		if _, ok := domain[value]; ok && len(domain) == 1 {
			delete(db.data, name)
		} else {
			delete(domain, value)
		}
	}
}

func (db *DB) GC() {
	now := time.Now()
	db.lock.Lock()
	defer db.lock.Unlock()
	for name, domain := range db.data {
		for value, expiresAt := range domain {
			if expiresAt.Before(now) {
				delete(domain, value)
			}
		}
		if len(domain) == 0 {
			delete(db.data, name)
		}
	}
}

func NewDB() *DB {
	db := &DB{
		data:    map[string]map[string]time.Time{},
		lock:    &sync.Mutex{},
		timeout: 2 * time.Minute,
	}
	go func() {
		for {
			time.Sleep(5 * db.timeout)
			db.GC()
		}
	}()
	return db
}