memoryweave/go/internal/distributed/distributed.go

// 织忆 MemoryWeave — 分布式同步 + CRDT + 速率限制
package distributed

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

// ─── 事件传播（HTTP 替代 Redis Streams）─────────────────

type Event struct {
	ID        string          `json:"id"`
	Type      string          `json:"type"`
	AgentID   string          `json:"agent_id"`
	Payload   json.RawMessage `json:"payload"`
	Timestamp time.Time       `json:"timestamp"`
}

type EventBus struct {
	mu        sync.RWMutex
	handlers  map[string]func(Event)
	peers     []string // 其他织忆实例的地址
	httpClient *http.Client
}

func NewEventBus(peers []string) *EventBus {
	return &EventBus{
		handlers:   make(map[string]func(Event)),
		peers:      peers,
		httpClient: &http.Client{Timeout: 5 * time.Second},
	}
}

// Subscribe 订阅事件
func (eb *EventBus) Subscribe(eventType string, handler func(Event)) {
	eb.mu.Lock()
	defer eb.mu.Unlock()
	eb.handlers[eventType] = handler
}

// Publish 发布本地事件并广播到所有 peer
func (eb *EventBus) Publish(event Event) {
	// 广播到 peers
	for _, peer := range eb.peers {
		go eb.sendToPeer(peer, event)
	}
	// 触发本地 handler
	eb.mu.RLock()
	if h, ok := eb.handlers[event.Type]; ok {
		go h(event)
	}
	eb.mu.RUnlock()
}

func (eb *EventBus) sendToPeer(peer string, event Event) {
	body, _ := json.Marshal(event)
	resp, err := eb.httpClient.Post(
		"http://"+peer+"/api/v1/events/receive",
		"application/json",
		bytes.NewReader(body))
	if err != nil {
		log.Printf("[eventbus] peer %s unreachable: %v", peer, err)
		return
	}
	resp.Body.Close()
}

// POST /api/v1/events/receive — 接收 peer 事件
func (eb *EventBus) ReceiveHandler(w http.ResponseWriter, r *http.Request) {
	var event Event
	if err := json.NewDecoder(r.Body).Decode(&event); err != nil {
		http.Error(w, "invalid event", 400)
		return
	}
	eb.mu.RLock()
	if h, ok := eb.handlers[event.Type]; ok {
		go h(event)
	}
	eb.mu.RUnlock()
	w.WriteHeader(200)
	json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

// ─── CRDT 合并 ──────────────────────────────────────────

type CRDTMerge struct {
	mu sync.RWMutex
}

// Merge 两条记录冲突时合并（时间戳 + 来源优先级）
func (c *CRDTMerge) Merge(a, b map[string]interface{}) map[string]interface{} {
	aTime := parseTime(a["updated_at"])
	bTime := parseTime(b["updated_at"])

	var winner, loser map[string]interface{}
	if aTime.After(bTime) {
		winner, loser = a, b
	} else {
		winner, loser = b, a
	}

	// 来源优先级：牧尘修正 > 时间戳
	if isMuchenCorrection(loser) {
		winner, loser = loser, winner
	}

	// 标记合并历史
	winner["merged_from"] = loser["id"]
	winner["merge_reason"] = "crdt_merge"
	winner["merged_at"] = time.Now().Format(time.RFC3339)

	return winner
}

func parseTime(v interface{}) time.Time {
	switch t := v.(type) {
	case time.Time:
		return t
	case string:
		parsed, _ := time.Parse(time.RFC3339, t)
		return parsed
	}
	return time.Now()
}

func isMuchenCorrection(record map[string]interface{}) bool {
	source, ok := record["source"].(string)
	return ok && (source == "muchen_oral" || source == "muchen_correction")
}

// ─── 速率限制（令牌桶）───────────────────────────────────

type RateLimiter struct {
	mu       sync.RWMutex
	buckets  map[string]*tokenBucket
	rate     float64 // 令牌/秒
	burst    int
}

type tokenBucket struct {
	tokens   float64
	lastTime time.Time
}

func NewRateLimiter(rate float64, burst int) *RateLimiter {
	return &RateLimiter{
		buckets: make(map[string]*tokenBucket),
		rate:    rate,
		burst:   burst,
	}
}

// Allow 检查指定 agent 是否允许本次请求
func (rl *RateLimiter) Allow(agentID string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	bucket, ok := rl.buckets[agentID]
	if !ok {
		bucket = &tokenBucket{tokens: float64(rl.burst), lastTime: time.Now()}
		rl.buckets[agentID] = bucket
	}

	// 补充令牌
	elapsed := time.Since(bucket.lastTime).Seconds()
	bucket.tokens += elapsed * rl.rate
	if bucket.tokens > float64(rl.burst) {
		bucket.tokens = float64(rl.burst)
	}
	bucket.lastTime = time.Now()

	if bucket.tokens >= 1 {
		bucket.tokens--
		return true
	}
	return false
}

// ─── 健康检查 ────────────────────────────────────────────

type HealthChecker struct {
	mu      sync.RWMutex
	started time.Time
	peers   map[string]time.Time // peer → last_seen
}

func NewHealthChecker(peers []string) *HealthChecker {
	h := &HealthChecker{
		started: time.Now(),
		peers:   make(map[string]time.Time),
	}
	for _, p := range peers {
		h.peers[p] = time.Time{}
	}
	return h
}

func (h *HealthChecker) Uptime() string { return time.Since(h.started).String() }

// ─── Prometheus 指标端点 ─────────────────────────────────

type MetricsHandler struct {
	dashboard interface{ Metrics() map[string]float64 }
	stats     func() map[string]interface{}
}

func NewMetricsHandler(dash interface{ Metrics() map[string]float64 }, statsFn func() map[string]interface{}) *MetricsHandler {
	return &MetricsHandler{dashboard: dash, stats: statsFn}
}

func (mh *MetricsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; version=0.0.4")

	metrics := mh.dashboard.Metrics()
	s := mh.stats()

	fmt.Fprintf(w, "# HELP zhiyi_recall_usefulness_rate Recall有用率\n")
	fmt.Fprintf(w, "# TYPE zhiyi_recall_usefulness_rate gauge\n")
	fmt.Fprintf(w, "zhiyi_recall_usefulness_rate %f\n", metrics["recall_usefulness_rate"])

	fmt.Fprintf(w, "# HELP zhiyi_recall_hit_rate Recall命中率\n")
	fmt.Fprintf(w, "# TYPE zhiyi_recall_hit_rate gauge\n")
	fmt.Fprintf(w, "zhiyi_recall_hit_rate %f\n", metrics["recall_hit_rate"])

	fmt.Fprintf(w, "# HELP zhiyi_gap_closure_rate 缺口闭环率\n")
	fmt.Fprintf(w, "# TYPE zhiyi_gap_closure_rate gauge\n")
	fmt.Fprintf(w, "zhiyi_gap_closure_rate %f\n", metrics["gap_closure_rate"])

	fmt.Fprintf(w, "# HELP zhiyi_episodes_total 总episode数\n")
	fmt.Fprintf(w, "# TYPE zhiyi_episodes_total gauge\n")
	if epCount, ok := s["episodes"]; ok {
		fmt.Fprintf(w, "zhiyi_episodes_total %v\n", epCount)
	}
}