1package pgs
  2
  3import (
  4	"context"
  5	"fmt"
  6	"io"
  7	"log/slog"
  8	"strings"
  9	"sync"
 10	"time"
 11
 12	"github.com/hashicorp/golang-lru/v2/expirable"
 13	"github.com/picosh/pico/pkg/shared"
 14	"github.com/picosh/utils/pipe"
 15)
 16
 17type PicoPubsub interface {
 18	io.Reader
 19	io.Writer
 20	io.Closer
 21}
 22
 23type PubsubPipe struct {
 24	Pipe *pipe.ReconnectReadWriteCloser
 25}
 26
 27func (p *PubsubPipe) Read(b []byte) (int, error) {
 28	return p.Pipe.Read(b)
 29}
 30
 31func (p *PubsubPipe) Write(b []byte) (int, error) {
 32	return p.Pipe.Write(b)
 33}
 34
 35func (p *PubsubPipe) Close() error {
 36	return p.Pipe.Close()
 37}
 38
 39func NewPubsubPipe(pipe *pipe.ReconnectReadWriteCloser) *PubsubPipe {
 40	return &PubsubPipe{
 41		Pipe: pipe,
 42	}
 43}
 44
 45type PubsubChan struct {
 46	Chan chan []byte
 47}
 48
 49func (p *PubsubChan) Read(b []byte) (int, error) {
 50	n := copy(b, <-p.Chan)
 51	return n, nil
 52}
 53
 54func (p *PubsubChan) Write(b []byte) (int, error) {
 55	p.Chan <- b
 56	return len(b), nil
 57}
 58
 59func (p *PubsubChan) Close() error {
 60	close(p.Chan)
 61	return nil
 62}
 63
 64func NewPubsubChan() *PubsubChan {
 65	return &PubsubChan{
 66		Chan: make(chan []byte),
 67	}
 68}
 69
 70func getSurrogateKey(userName, projectName string) string {
 71	return fmt.Sprintf("%s-%s", userName, projectName)
 72}
 73
 74func CreatePubCacheDrain(ctx context.Context, logger *slog.Logger) *pipe.ReconnectReadWriteCloser {
 75	info := shared.NewPicoPipeClient()
 76	send := pipe.NewReconnectReadWriteCloser(
 77		ctx,
 78		logger,
 79		info,
 80		"pub to cache-drain",
 81		"pub cache-drain -b=false",
 82		100,
 83		-1,
 84	)
 85	return send
 86}
 87
 88func CreateSubCacheDrain(ctx context.Context, logger *slog.Logger) *pipe.ReconnectReadWriteCloser {
 89	info := shared.NewPicoPipeClient()
 90	send := pipe.NewReconnectReadWriteCloser(
 91		ctx,
 92		logger,
 93		info,
 94		"sub to cache-drain",
 95		"sub cache-drain -k",
 96		100,
 97		-1,
 98	)
 99	return send
100}
101
102// purgeCache send a pipe pub to the pgs web instance which purges
103// cached entries for a given subdomain (like "fakeuser-www-proj"). We set a
104// "surrogate-key: <subdomain>" header on every pgs response which ensures all
105// cached assets for a given subdomain are grouped under a single key (which is
106// separate from the "GET-https-example.com-/path" key used for serving files
107// from the cache).
108func purgeCache(cfg *Config, writer io.Writer, surrogate string) error {
109	cfg.Logger.Info("purging cache", "surrogate", surrogate)
110	_, err := writer.Write([]byte(surrogate + "\n"))
111	return err
112}
113
114func purgeAllCache(cfg *Config, writer io.Writer) error {
115	return purgeCache(cfg, writer, "*")
116}
117
118// SiteCache holds the http cache's responses and indexes their keys by site,
119// so purging one site doesn't scan the whole cache.
120type SiteCache struct {
121	*expirable.LRU[string, []byte]
122	mu    sync.Mutex
123	sites map[string]map[string]struct{}
124}
125
126func NewSiteCache(size int, ttl time.Duration, onEvict func(key string, value []byte)) *SiteCache {
127	c := &SiteCache{sites: map[string]map[string]struct{}{}}
128	c.LRU = expirable.NewLRU(size, func(key string, value []byte) {
129		c.unindex(key)
130		if onEvict != nil {
131			onEvict(key, value)
132		}
133	}, ttl)
134	return c
135}
136
137// cacheKeySite returns the site a PgsCacheKey belongs to.
138func cacheKeySite(key string) string {
139	site, _, _ := strings.Cut(key, "__")
140	return site
141}
142
143func (c *SiteCache) Add(key string, value []byte) bool {
144	evicted := c.LRU.Add(key, value)
145	site := cacheKeySite(key)
146	c.mu.Lock()
147	defer c.mu.Unlock()
148	if c.sites[site] == nil {
149		c.sites[site] = map[string]struct{}{}
150	}
151	c.sites[site][key] = struct{}{}
152	return evicted
153}
154
155// PurgeSite removes every cached response for a site. The index lock is
156// released before touching the LRU, whose eviction callback takes it.
157func (c *SiteCache) PurgeSite(site string) {
158	c.mu.Lock()
159	keys := c.sites[site]
160	delete(c.sites, site)
161	c.mu.Unlock()
162	for key := range keys {
163		c.Remove(key)
164	}
165}
166
167func (c *SiteCache) unindex(key string) {
168	site := cacheKeySite(key)
169	c.mu.Lock()
170	defer c.mu.Unlock()
171	if keys, ok := c.sites[site]; ok {
172		delete(keys, key)
173		if len(keys) == 0 {
174			delete(c.sites, site)
175		}
176	}
177}