468 lines
11 KiB
Go
468 lines
11 KiB
Go
// 织忆 MemoryWeave — Redis 原生客户端（RESP 协议，零外部依赖）
|
||
package storage
|
||
|
||
import (
|
||
"bufio"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net"
|
||
"os"
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// ─── RESP 客户端 ───────────────────────────────────────
|
||
|
||
// RedisConn 原生 TCP 连接 + RESP 读写
|
||
type RedisConn struct {
|
||
mu sync.Mutex
|
||
conn net.Conn
|
||
r *bufio.Reader
|
||
}
|
||
|
||
func DialRedis(addr string) (*RedisConn, error) {
|
||
conn, err := net.DialTimeout("tcp", addr, 3*time.Second)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
return &RedisConn{conn: conn, r: bufio.NewReader(conn)}, nil
|
||
}
|
||
|
||
func (rc *RedisConn) Close() error { return rc.conn.Close() }
|
||
|
||
// Do 执行 Redis 命令，返回 RESP 解析结果
|
||
func (rc *RedisConn) Do(args ...string) (interface{}, error) {
|
||
rc.mu.Lock()
|
||
defer rc.mu.Unlock()
|
||
|
||
// 编码 RESP
|
||
cmd := fmt.Sprintf("*%d\r\n", len(args))
|
||
for _, a := range args {
|
||
cmd += fmt.Sprintf("$%d\r\n%s\r\n", len(a), a)
|
||
}
|
||
|
||
if _, err := rc.conn.Write([]byte(cmd)); err != nil {
|
||
return nil, fmt.Errorf("redis write: %w", err)
|
||
}
|
||
return rc.readRESP()
|
||
}
|
||
|
||
func (rc *RedisConn) readRESP() (interface{}, error) {
|
||
line, err := rc.r.ReadString('\n')
|
||
if err != nil {
|
||
return nil, fmt.Errorf("redis read: %w", err)
|
||
}
|
||
line = strings.TrimSuffix(line, "\r\n")
|
||
|
||
switch {
|
||
case strings.HasPrefix(line, "+"):
|
||
return line[1:], nil
|
||
case strings.HasPrefix(line, "-"):
|
||
return nil, errors.New(line[1:])
|
||
case strings.HasPrefix(line, ":"):
|
||
return strconv.ParseInt(line[1:], 10, 64)
|
||
case strings.HasPrefix(line, "$"):
|
||
length, _ := strconv.Atoi(line[1:])
|
||
if length < 0 {
|
||
return nil, nil
|
||
}
|
||
buf := make([]byte, length+2)
|
||
if _, err := io.ReadFull(rc.r, buf); err != nil {
|
||
return nil, err
|
||
}
|
||
return string(buf[:length]), nil
|
||
case strings.HasPrefix(line, "*"):
|
||
count, _ := strconv.Atoi(line[1:])
|
||
if count < 0 {
|
||
return nil, nil
|
||
}
|
||
arr := make([]interface{}, count)
|
||
for i := 0; i < count; i++ {
|
||
arr[i], _ = rc.readRESP()
|
||
}
|
||
return arr, nil
|
||
}
|
||
return line, nil
|
||
}
|
||
|
||
// ─── Redis 客户端封装 ─────────────────────────────────
|
||
|
||
type RedisClient struct {
|
||
conn *RedisConn
|
||
addr string
|
||
}
|
||
|
||
func NewRedisClient() *RedisClient {
|
||
addr := os.Getenv("REDIS_ADDR")
|
||
if addr == "" {
|
||
addr = "127.0.0.1:6379"
|
||
}
|
||
return &RedisClient{addr: addr}
|
||
}
|
||
|
||
func (rc *RedisClient) Connect() error {
|
||
if rc.conn != nil {
|
||
return nil
|
||
}
|
||
conn, err := DialRedis(rc.addr)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
rc.conn = conn
|
||
// 健康检查
|
||
if _, err := conn.Do("PING"); err != nil {
|
||
conn.Close()
|
||
rc.conn = nil
|
||
return err
|
||
}
|
||
log.Printf("[redis] connected to %s", rc.addr)
|
||
return nil
|
||
}
|
||
|
||
// ─── 基础命令 ─────────────────────────────────────────
|
||
|
||
func (rc *RedisClient) Set(key, value string, ttl time.Duration) error {
|
||
if rc.conn == nil {
|
||
return errors.New("redis: not connected")
|
||
}
|
||
args := []string{"SET", key, value}
|
||
if ttl > 0 {
|
||
args = append(args, "EX", strconv.Itoa(int(ttl.Seconds())))
|
||
}
|
||
_, err := rc.conn.Do(args...)
|
||
return err
|
||
}
|
||
|
||
func (rc *RedisClient) Get(key string) (string, error) {
|
||
if rc.conn == nil {
|
||
return "", errors.New("redis: not connected")
|
||
}
|
||
v, err := rc.conn.Do("GET", key)
|
||
if err != nil {
|
||
return "", err
|
||
}
|
||
if v == nil {
|
||
return "", nil
|
||
}
|
||
return v.(string), nil
|
||
}
|
||
|
||
func (rc *RedisClient) Del(keys ...string) error {
|
||
if rc.conn == nil {
|
||
return errors.New("redis: not connected")
|
||
}
|
||
args := append([]string{"DEL"}, keys...)
|
||
_, err := rc.conn.Do(args...)
|
||
return err
|
||
}
|
||
|
||
func (rc *RedisClient) Expire(key string, ttl time.Duration) error {
|
||
if rc.conn == nil {
|
||
return errors.New("redis: not connected")
|
||
}
|
||
_, err := rc.conn.Do("EXPIRE", key, strconv.Itoa(int(ttl.Seconds())))
|
||
return err
|
||
}
|
||
|
||
func (rc *RedisClient) Incr(key string) (int64, error) {
|
||
if rc.conn == nil {
|
||
return 0, errors.New("redis: not connected")
|
||
}
|
||
v, err := rc.conn.Do("INCR", key)
|
||
if err != nil {
|
||
return 0, err
|
||
}
|
||
return v.(int64), nil
|
||
}
|
||
|
||
// ─── Hash ─────────────────────────────────────────────
|
||
|
||
func (rc *RedisClient) HSet(key, field, value string) error {
|
||
if rc.conn == nil {
|
||
return errors.New("redis: not connected")
|
||
}
|
||
_, err := rc.conn.Do("HSET", key, field, value)
|
||
return err
|
||
}
|
||
|
||
// HDel 从 Hash 中删除一个或多个 field
|
||
func (rc *RedisClient) HDel(key, field string) error {
|
||
if rc.conn == nil {
|
||
return errors.New("redis: not connected")
|
||
}
|
||
_, err := rc.conn.Do("HDEL", key, field)
|
||
return err
|
||
}
|
||
|
||
func (rc *RedisClient) HGetAll(key string) (map[string]string, error) {
|
||
if rc.conn == nil {
|
||
return nil, errors.New("redis: not connected")
|
||
}
|
||
v, err := rc.conn.Do("HGETALL", key)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
arr, ok := v.([]interface{})
|
||
if !ok {
|
||
return nil, errors.New("redis: unexpected HGETALL response")
|
||
}
|
||
result := make(map[string]string)
|
||
for i := 0; i < len(arr); i += 2 {
|
||
if s, ok := arr[i].(string); ok {
|
||
if val, ok := arr[i+1].(string); ok {
|
||
result[s] = val
|
||
}
|
||
}
|
||
}
|
||
return result, nil
|
||
}
|
||
|
||
// ─── Pub/Sub ──────────────────────────────────────────
|
||
|
||
func (rc *RedisClient) Publish(channel, message string) error {
|
||
if rc.conn == nil {
|
||
return errors.New("redis: not connected")
|
||
}
|
||
_, err := rc.conn.Do("PUBLISH", channel, message)
|
||
return err
|
||
}
|
||
|
||
// ─── 全局客户端 ───────────────────────────────────────
|
||
|
||
var redisClient *RedisClient
|
||
var redisOnce sync.Once
|
||
|
||
func GetRedisClient() *RedisClient {
|
||
redisOnce.Do(func() {
|
||
redisClient = NewRedisClient()
|
||
if err := redisClient.Connect(); err != nil {
|
||
log.Printf("[redis] connect failed: %v — falling back to in-memory", err)
|
||
redisClient = nil
|
||
}
|
||
})
|
||
return redisClient
|
||
}
|
||
|
||
// ─── Redis 限流器 ─────────────────────────────────────
|
||
|
||
type RedisRateLimiter struct {
|
||
client *RedisClient
|
||
mu sync.Mutex
|
||
local map[string]*tokenBucket // fallback
|
||
}
|
||
|
||
type tokenBucket struct {
|
||
tokens float64
|
||
lastTime time.Time
|
||
rate float64
|
||
burst float64
|
||
}
|
||
|
||
func NewRedisRateLimiter() *RedisRateLimiter {
|
||
return &RedisRateLimiter{
|
||
client: GetRedisClient(),
|
||
local: make(map[string]*tokenBucket),
|
||
}
|
||
}
|
||
|
||
func (rl *RedisRateLimiter) Allow(key string, rate, burst float64) bool {
|
||
if rl.client != nil {
|
||
// 用 Redis INCR + EXPIRE 做简单计数限流
|
||
redisKey := "zhiyi:ratelimit:" + key
|
||
count, err := rl.client.Incr(redisKey)
|
||
if err == nil && count == 1 {
|
||
rl.client.Expire(redisKey, time.Minute)
|
||
}
|
||
return count <= int64(burst)
|
||
}
|
||
|
||
// Fallback: 内存令牌桶
|
||
rl.mu.Lock()
|
||
defer rl.mu.Unlock()
|
||
bucket, ok := rl.local[key]
|
||
if !ok {
|
||
bucket = &tokenBucket{tokens: burst, lastTime: time.Now(), rate: rate, burst: burst}
|
||
rl.local[key] = bucket
|
||
}
|
||
now := time.Now()
|
||
elapsed := now.Sub(bucket.lastTime).Seconds()
|
||
bucket.tokens += elapsed * bucket.rate
|
||
if bucket.tokens > bucket.burst {
|
||
bucket.tokens = bucket.burst
|
||
}
|
||
bucket.lastTime = now
|
||
if bucket.tokens >= 1 {
|
||
bucket.tokens--
|
||
return true
|
||
}
|
||
return false
|
||
}
|
||
|
||
// ─── Redis 缓存 ───────────────────────────────────────
|
||
|
||
type RedisCache struct {
|
||
client *RedisClient
|
||
ttl time.Duration
|
||
mu sync.RWMutex
|
||
fallback map[string]*cachedItem
|
||
}
|
||
|
||
type cachedItem struct {
|
||
data []byte
|
||
expiresAt time.Time
|
||
}
|
||
|
||
func NewRedisCache(ttl time.Duration) *RedisCache {
|
||
return &RedisCache{
|
||
client: GetRedisClient(),
|
||
ttl: ttl,
|
||
fallback: make(map[string]*cachedItem),
|
||
}
|
||
}
|
||
|
||
func (rc *RedisCache) Get(key string) ([]byte, bool) {
|
||
if rc.client != nil {
|
||
v, err := rc.client.Get("zhiyi:cache:" + key)
|
||
if err == nil && v != "" {
|
||
return []byte(v), true
|
||
}
|
||
}
|
||
|
||
rc.mu.RLock()
|
||
item, ok := rc.fallback[key]
|
||
rc.mu.RUnlock()
|
||
if !ok || time.Now().After(item.expiresAt) {
|
||
return nil, false
|
||
}
|
||
return item.data, true
|
||
}
|
||
|
||
func (rc *RedisCache) Set(key string, data []byte) {
|
||
if rc.client != nil {
|
||
rc.client.Set("zhiyi:cache:"+key, string(data), rc.ttl)
|
||
return
|
||
}
|
||
rc.mu.Lock()
|
||
rc.fallback[key] = &cachedItem{data: data, expiresAt: time.Now().Add(rc.ttl)}
|
||
rc.mu.Unlock()
|
||
}
|
||
|
||
func (rc *RedisCache) Del(key string) {
|
||
if rc.client != nil {
|
||
rc.client.Del("zhiyi:cache:" + key)
|
||
return
|
||
}
|
||
rc.mu.Lock()
|
||
delete(rc.fallback, key)
|
||
rc.mu.Unlock()
|
||
}
|
||
|
||
// ─── Redis 心跳 ───────────────────────────────────────
|
||
|
||
type Heartbeat struct {
|
||
client *RedisClient
|
||
instance string
|
||
done chan struct{}
|
||
}
|
||
|
||
func NewHeartbeat(instance string) *Heartbeat {
|
||
h := &Heartbeat{
|
||
client: GetRedisClient(),
|
||
instance: instance,
|
||
done: make(chan struct{}),
|
||
}
|
||
go h.beat()
|
||
return h
|
||
}
|
||
|
||
func (h *Heartbeat) beat() {
|
||
ticker := time.NewTicker(10 * time.Second)
|
||
defer ticker.Stop()
|
||
for {
|
||
select {
|
||
case <-ticker.C:
|
||
if h.client != nil {
|
||
h.client.Set("zhiyi:heartbeat:"+h.instance, time.Now().Format(time.RFC3339), 30*time.Second)
|
||
}
|
||
case <-h.done:
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
func (h *Heartbeat) Stop() { close(h.done) }
|
||
|
||
// ─── Redis 事件总线 ───────────────────────────────────
|
||
|
||
type RedisEventBus struct {
|
||
client *RedisClient
|
||
}
|
||
|
||
func NewRedisEventBus() *RedisEventBus {
|
||
return &RedisEventBus{client: GetRedisClient()}
|
||
}
|
||
|
||
func (reb *RedisEventBus) Publish(channel string, data []byte) error {
|
||
if reb.client != nil {
|
||
return reb.client.Publish(channel, string(data))
|
||
}
|
||
return errors.New("redis: not available")
|
||
}
|
||
|
||
// ─── 自优化指标存储 ──────────────────────────────────
|
||
|
||
type MetricsStore struct {
|
||
client *RedisClient
|
||
mu sync.RWMutex
|
||
local map[string]map[string]float64
|
||
}
|
||
|
||
var GlobalMetricsStore = &MetricsStore{
|
||
client: GetRedisClient(),
|
||
local: make(map[string]map[string]float64),
|
||
}
|
||
|
||
func (ms *MetricsStore) Set(date, metric string, value float64) {
|
||
if ms.client != nil {
|
||
ms.client.HSet("zhiyi:metrics:"+date, metric, fmt.Sprintf("%.4f", value))
|
||
ms.client.Expire("zhiyi:metrics:"+date, 90*24*time.Hour)
|
||
return
|
||
}
|
||
ms.mu.Lock()
|
||
if ms.local[date] == nil {
|
||
ms.local[date] = make(map[string]float64)
|
||
}
|
||
ms.local[date][metric] = value
|
||
ms.mu.Unlock()
|
||
}
|
||
|
||
func (ms *MetricsStore) GetRange(days int) []map[string]interface{} {
|
||
var result []map[string]interface{}
|
||
for t := time.Now(); t.After(time.Now().AddDate(0, 0, -days)); t = t.AddDate(0, 0, -1) {
|
||
date := t.Format("2006-01-02")
|
||
if ms.client != nil {
|
||
data, err := ms.client.HGetAll("zhiyi:metrics:" + date)
|
||
if err == nil && len(data) > 0 {
|
||
result = append(result, map[string]interface{}{
|
||
"date": date,
|
||
"metrics": data,
|
||
})
|
||
}
|
||
} else {
|
||
ms.mu.RLock()
|
||
if m, ok := ms.local[date]; ok {
|
||
result = append(result, map[string]interface{}{
|
||
"date": date,
|
||
"metrics": m,
|
||
})
|
||
}
|
||
ms.mu.RUnlock()
|
||
}
|
||
}
|
||
return result
|
||
}
|