init

internal/tenant/pool.go

@@ -0,0 +1,176 @@
+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)
+}