清理代码

pkg/middleware/cache/cache.go

@@ -0,0 +1,48 @@
+package cache
+
+import (
+	"io"
+	"sync"
+	"time"
+
+	"github.com/pkg/errors"
+)
+
+type CacheContent struct {
+	io.ReadSeekCloser
+	Length       int
+	LastModified time.Time
+}
+
+func (c *CacheContent) ReadToString() (string, error) {
+	all, err := io.ReadAll(c)
+	if err != nil {
+		return "", err
+	}
+	return string(all), nil
+}
+
+type Cache interface {
+	Put(key string, reader io.Reader) error
+	// Get return CacheContent or nil when put nil io.reader
+	Get(key string) (*CacheContent, error)
+	Delete(pattern string) error
+	io.Closer
+}
+
+var ErrCacheOutOfMemory = errors.New("内容无法被缓存，超过最大限定值")
+
+// TODO: 优化锁结构
+// 复杂场景请使用其他缓存服务
+
+type CacheMemory struct {
+	l          sync.RWMutex
+	data       map[string]*[]byte
+	lastModify map[string]time.Time
+	sizeGlobal int
+	sizeItem   int
+
+	current int
+	cache   []byte
+	ordered []string
+}

pkg/middleware/cache/cache_memory.go

@@ -0,0 +1,139 @@
+package cache
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	"gopkg.d7z.net/gitea-pages/pkg/utils"
+)
+
+func NewCacheMemory(maxUsage, maxGlobalUsage int) *CacheMemory {
+	return &CacheMemory{
+		data:       make(map[string]*[]byte),
+		lastModify: make(map[string]time.Time),
+		l:          sync.RWMutex{},
+		sizeGlobal: maxGlobalUsage,
+		sizeItem:   maxUsage,
+
+		cache:   make([]byte, maxUsage+1),
+		ordered: make([]string, 0),
+	}
+}
+
+func (c *CacheMemory) Put(key string, reader io.Reader) error {
+	c.l.Lock()
+	defer c.l.Unlock()
+	size := 0
+	// 可以指定空的 reader 作为 404 缓存
+	if reader != nil {
+		var err error
+		size, err = io.ReadAtLeast(reader, c.cache, 1)
+		if err != nil {
+			return err
+		}
+	}
+	if size == len(c.cache) {
+		return ErrCacheOutOfMemory
+	}
+	currentItemSize := 0
+	if data, ok := c.data[key]; ok {
+		currentItemSize = len(*data)
+	}
+	available := c.sizeGlobal + currentItemSize - (c.current + size)
+	if available < 0 {
+		// 清理旧的内容
+		count := 0
+		for i, k := range c.ordered {
+			available += len(*c.data[k])
+			if available > 0 {
+				break
+			}
+			count = i + 1
+		}
+
+		if available < 0 {
+			// 清理全部内容也无法留出空间
+			return ErrCacheOutOfMemory
+		}
+		for _, s := range c.ordered[:count] {
+			delete(c.data, s)
+			delete(c.lastModify, s)
+		}
+		c.ordered = c.ordered[count:]
+	}
+
+	if reader != nil {
+		dest := make([]byte, size)
+		copy(dest, c.cache[:size])
+		c.data[key] = &dest
+		c.lastModify[key] = time.Now()
+
+		c.current -= currentItemSize
+		c.current += len(dest)
+	} else {
+		c.data[key] = nil
+		c.lastModify[key] = time.Now()
+
+		c.current -= currentItemSize
+	}
+
+	nextOrdered := make([]string, 0, len(c.ordered))
+	for _, s := range c.ordered {
+		if s != key {
+			nextOrdered = append(nextOrdered, s)
+		}
+	}
+	c.ordered = append(nextOrdered, key)
+	return nil
+}
+
+func (c *CacheMemory) Get(key string) (*CacheContent, error) {
+	c.l.RLock()
+	defer c.l.RUnlock()
+	if i, ok := c.data[key]; ok {
+		if i == nil {
+			return nil, nil
+		}
+
+		return &CacheContent{
+			ReadSeekCloser: utils.NopCloser{
+				bytes.NewReader(*i),
+			},
+			Length:       len(*i),
+			LastModified: c.lastModify[key],
+		}, nil
+	}
+	return nil, os.ErrNotExist
+}
+
+func (c *CacheMemory) Delete(pattern string) error {
+	c.l.Lock()
+	defer c.l.Unlock()
+	nextOrder := make([]string, 0, len(c.ordered))
+	for _, key := range c.ordered {
+		if strings.HasPrefix(key, pattern) {
+			c.current -= len(*c.data[key])
+			delete(c.data, key)
+			delete(c.lastModify, key)
+		} else {
+			nextOrder = append(nextOrder, key)
+		}
+	}
+	clear(c.ordered)
+	c.ordered = nextOrder
+	return nil
+}
+
+func (c *CacheMemory) Close() error {
+	c.l.Lock()
+	defer c.l.Unlock()
+	clear(c.ordered)
+	clear(c.data)
+	clear(c.lastModify)
+	c.current = 0
+	return nil
+}

pkg/middleware/cache/cache_test.go

@@ -0,0 +1,75 @@
+package cache
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCacheGetPutDelete(t *testing.T) {
+	memory := NewCacheMemory(1024, 10240)
+
+	require.NoError(t, memory.Put("hello", strings.NewReader("world")))
+
+	value, err := memory.Get("hello")
+	require.NoError(t, err)
+	all, err := io.ReadAll(value)
+	require.NoError(t, err)
+	require.Equal(t, "world", string(all))
+	require.Equal(t, 5, memory.current)
+
+	require.NoError(t, memory.Put("hello", strings.NewReader("kotlin")))
+
+	value, err = memory.Get("hello")
+	require.NoError(t, err)
+	all, err = io.ReadAll(value)
+	require.NoError(t, err)
+	require.Equal(t, "kotlin", string(all))
+	require.Equal(t, 6, memory.current)
+	require.Equal(t, 1, len(memory.data))
+
+	require.NoError(t, memory.Put("data", strings.NewReader("kotlin")))
+	require.Equal(t, 12, memory.current)
+	require.Equal(t, 2, len(memory.data))
+	require.Equal(t, 2, len(memory.ordered))
+
+	require.NoError(t, memory.Delete("hello"))
+	value, err = memory.Get("hello")
+	require.Error(t, err)
+	require.Equal(t, 1, len(memory.data))
+	require.Equal(t, 1, len(memory.ordered))
+
+	require.NoError(t, memory.Put("hello", nil))
+	value, err = memory.Get("hello")
+	require.NoError(t, err)
+	require.Nil(t, value)
+}
+
+func TestCacheLimit(t *testing.T) {
+	memory := NewCacheMemory(5, 5*5)
+	require.NoError(t, memory.Put("hello", strings.NewReader("world")))
+	require.Equal(t, 5, memory.current)
+	require.ErrorIs(t, memory.Put("hello", strings.NewReader("world1")), ErrCacheOutOfMemory)
+	require.Equal(t, 5, memory.current)
+	for i := 0; i < 4; i++ {
+		require.NoError(t, memory.Put(fmt.Sprintf("hello-%d", i), strings.NewReader("govet")))
+	}
+	value, err := memory.Get("hello")
+	require.NoError(t, err)
+	all, err := io.ReadAll(value)
+	require.NoError(t, err)
+	require.Equal(t, "world", string(all))
+
+	require.NoError(t, memory.Put("test", strings.NewReader("govet")))
+
+	value, err = memory.Get("hello")
+	require.ErrorIs(t, err, os.ErrNotExist)
+
+	require.Equal(t, 5, len(memory.data))
+	require.Equal(t, 5, len(memory.ordered))
+	require.Equal(t, 5, len(memory.lastModify))
+}

pkg/middleware/config/config.go

@@ -0,0 +1,55 @@
+package config
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/url"
+	"strconv"
+	"strings"
+	"time"
+)
+
+const TtlKeep = -1
+
+type KVConfig interface {
+	Put(ctx context.Context, key string, value string, ttl time.Duration) error
+	Get(ctx context.Context, key string) (string, error)
+	Delete(ctx context.Context, key string) error
+	io.Closer
+}
+
+func NewAutoConfig(src string) (KVConfig, error) {
+	if src == "" ||
+		strings.HasPrefix(src, "./") ||
+		strings.HasPrefix(src, "/") ||
+		strings.HasPrefix(src, "\\") ||
+		strings.HasPrefix(src, ".\\") {
+		return NewConfigMemory(src)
+	}
+	parse, err := url.Parse(src)
+	if err != nil {
+		return nil, err
+	}
+	switch parse.Scheme {
+	case "local":
+		return NewConfigMemory(parse.Path)
+	case "redis":
+		query := parse.Query()
+		pass := query.Get("pass")
+		if pass == "" {
+			pass = query.Get("password")
+		}
+		db := strings.TrimPrefix(parse.Path, "/")
+		if db == "" {
+			db = "0"
+		}
+		dbi, err := strconv.Atoi(db)
+		if err != nil {
+			return nil, err
+		}
+		return NewConfigRedis(parse.Host, pass, dbi)
+	default:
+		return nil, fmt.Errorf("unsupported scheme: %s", parse.Scheme)
+	}
+}

pkg/middleware/config/config_memory.go

@@ -0,0 +1,106 @@
+package config
+
+import (
+	"context"
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"sync"
+	"time"
+
+	"go.uber.org/zap"
+)
+
+// ConfigMemory 一个简单的内存配置归档，仅用于测试
+type ConfigMemory struct {
+	data  sync.Map
+	store string
+}
+
+func NewConfigMemory(store string) (KVConfig, error) {
+	ret := &ConfigMemory{
+		store: store,
+		data:  sync.Map{},
+	}
+	if store != "" {
+		zap.L().Info("parse config from store", zap.String("store", store))
+		if err := os.MkdirAll(filepath.Dir(store), 0o755); err != nil && !os.IsExist(err) {
+			return nil, err
+		}
+		item := make(map[string]ConfigContent)
+		data, err := os.ReadFile(store)
+		if err != nil && !os.IsNotExist(err) {
+			return nil, err
+		}
+		if err == nil {
+			err = json.Unmarshal(data, &item)
+			if err != nil {
+				return nil, err
+			}
+		}
+		for key, content := range item {
+			if content.Ttl == nil || time.Now().Before(*content.Ttl) {
+				ret.data.Store(key, content)
+			}
+		}
+		clear(item)
+	}
+	return ret, nil
+}
+
+type ConfigContent struct {
+	Data string     `json:"data"`
+	Ttl  *time.Time `json:"ttl,omitempty"`
+}
+
+func (m *ConfigMemory) Put(ctx context.Context, key string, value string, ttl time.Duration) error {
+	d := time.Now().Add(ttl)
+	td := &d
+	if ttl == -1 {
+		td = nil
+	}
+	m.data.Store(key, ConfigContent{
+		Data: value,
+		Ttl:  td,
+	})
+	return nil
+}
+
+func (m *ConfigMemory) Get(ctx context.Context, key string) (string, error) {
+	if value, ok := m.data.Load(key); ok {
+		content := value.(ConfigContent)
+		if content.Ttl != nil && time.Now().After(*content.Ttl) {
+			return "", os.ErrNotExist
+		}
+		return content.Data, nil
+	}
+	return "", os.ErrNotExist
+}
+
+func (m *ConfigMemory) Delete(ctx context.Context, key string) error {
+	m.data.Delete(key)
+	return nil
+}
+
+func (m *ConfigMemory) Close() error {
+	defer m.data.Clear()
+	if m.store != "" {
+		item := make(map[string]ConfigContent)
+		now := time.Now()
+		m.data.Range(
+			func(key, value interface{}) bool {
+				content := value.(ConfigContent)
+				if content.Ttl == nil || now.Before(*content.Ttl) {
+					item[key.(string)] = content
+				}
+				return true
+			})
+		zap.L().Debug("回写内容到本地存储", zap.String("store", m.store), zap.Int("length", len(item)))
+		saved, err := json.Marshal(item)
+		if err != nil {
+			return err
+		}
+		return os.WriteFile(m.store, saved, 0o600)
+	}
+	return nil
+}

pkg/middleware/config/config_redis.go

@@ -0,0 +1,60 @@
+package config
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/valkey-io/valkey-go"
+
+	"go.uber.org/zap"
+)
+
+type ConfigRedis struct {
+	client valkey.Client
+}
+
+func NewConfigRedis(addr string, password string, db int) (*ConfigRedis, error) {
+	if addr == "" {
+		return nil, fmt.Errorf("addr is empty")
+	}
+	zap.L().Debug("connect redis", zap.String("addr", addr))
+	client, err := valkey.NewClient(valkey.ClientOption{
+		InitAddress: []string{addr},
+		Password:    password,
+		SelectDB:    db,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return &ConfigRedis{
+		client: client,
+	}, nil
+}
+
+func (r *ConfigRedis) Put(ctx context.Context, key string, value string, ttl time.Duration) error {
+	builder := r.client.B().Set().Key(key).Value(value)
+	if ttl != TtlKeep {
+		builder.Ex(ttl)
+	}
+	return r.client.Do(ctx, builder.Build()).Error()
+}
+
+func (r *ConfigRedis) Get(ctx context.Context, key string) (string, error) {
+	v, err := r.client.Do(ctx, r.client.B().Get().Key(key).Build()).ToString()
+	if err != nil && errors.Is(err, valkey.Nil) {
+		return "", os.ErrNotExist
+	}
+	return v, err
+}
+
+func (r *ConfigRedis) Delete(ctx context.Context, key string) error {
+	return r.client.Do(ctx, r.client.B().Del().Key(key).Build()).Error()
+}
+
+func (r *ConfigRedis) Close() error {
+	r.client.Close()
+	return nil
+}