writekit/internal/tenant/pool.go

package tenant

import (
	"container/list"
	"database/sql"
	"sync"
	"time"
)

const (
	maxOpenConns     = 500
	cacheTTL         = 5 * time.Minute
	cacheCleanupFreq = time.Minute
)

type conn struct {
	db       *sql.DB
	tenantID string
}

// Pool manages SQLite connections for tenants with LRU eviction.
type Pool struct {
	dataDir  string
	mu       sync.Mutex
	conns    map[string]*list.Element
	lru      *list.List
	inMemory map[string]bool
}

func NewPool(dataDir string) *Pool {
	return &Pool{
		dataDir:  dataDir,
		conns:    make(map[string]*list.Element),
		lru:      list.New(),
		inMemory: make(map[string]bool),
	}
}

func (p *Pool) MarkAsDemo(tenantID string) {
	p.mu.Lock()
	defer p.mu.Unlock()
	p.inMemory[tenantID] = true
}

func (p *Pool) Get(tenantID string) (*sql.DB, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if elem, ok := p.conns[tenantID]; ok {
		p.lru.MoveToFront(elem)
		return elem.Value.(*conn).db, nil
	}

	useInMemory := p.inMemory[tenantID]
	db, err := openDB(p.dataDir, tenantID, useInMemory)
	if err != nil {
		return nil, err
	}

	for p.lru.Len() >= maxOpenConns {
		p.evictOldest()
	}

	c := &conn{db: db, tenantID: tenantID}
	elem := p.lru.PushFront(c)
	p.conns[tenantID] = elem

	return db, nil
}

func (p *Pool) evictOldest() {
	elem := p.lru.Back()
	if elem == nil {
		return
	}
	c := elem.Value.(*conn)
	c.db.Close()
	delete(p.conns, c.tenantID)
	delete(p.inMemory, c.tenantID)
	p.lru.Remove(elem)
}

func (p *Pool) Evict(tenantID string) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if elem, ok := p.conns[tenantID]; ok {
		c := elem.Value.(*conn)
		c.db.Close()
		delete(p.conns, tenantID)
		p.lru.Remove(elem)
	}
	delete(p.inMemory, tenantID)
}

func (p *Pool) Close() {
	p.mu.Lock()
	defer p.mu.Unlock()

	for p.lru.Len() > 0 {
		p.evictOldest()
	}
}

type cacheEntry struct {
	tenantID  string
	expiresAt time.Time
}

// Cache stores subdomain to tenant ID mappings.
type Cache struct {
	mu    sync.RWMutex
	items map[string]cacheEntry
	stop  chan struct{}
}

func NewCache() *Cache {
	c := &Cache{
		items: make(map[string]cacheEntry),
		stop:  make(chan struct{}),
	}
	go c.cleanup()
	return c
}

func (c *Cache) Get(subdomain string) (string, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	entry, ok := c.items[subdomain]
	if !ok || time.Now().After(entry.expiresAt) {
		return "", false
	}
	return entry.tenantID, true
}

func (c *Cache) Set(subdomain, tenantID string) {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.items[subdomain] = cacheEntry{
		tenantID:  tenantID,
		expiresAt: time.Now().Add(cacheTTL),
	}
}

func (c *Cache) Delete(subdomain string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	delete(c.items, subdomain)
}

func (c *Cache) cleanup() {
	ticker := time.NewTicker(cacheCleanupFreq)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			c.mu.Lock()
			now := time.Now()
			for k, v := range c.items {
				if now.After(v.expiresAt) {
					delete(c.items, k)
				}
			}
			c.mu.Unlock()
		case <-c.stop:
			return
		}
	}
}

func (c *Cache) Close() {
	close(c.stop)
}