Agent-System/cmd/events/main.go

package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"chat-agent-system/internal/config"
	"chat-agent-system/internal/db"
	"chat-agent-system/internal/events"

	"github.com/jackc/pgx/v5/pgxpool"
)

func main() {
	envPath := flag.String("env", ".env", "环境配置文件路径")
	flag.Parse()

	cfg, err := config.Load(*envPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "events: load config: %v\n", os.ErrInvalid)
		os.Exit(1)
	}
	port := cfg.HTTPPort
	if v := os.Getenv("EVENTS_HTTP_PORT"); v != "" {
		if p, perr := parsePort(v); perr == nil {
			port = p
		}
	}

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	pc, err := pgxpool.ParseConfig(cfg.ContextDSN())
	if err != nil {
		fmt.Fprintf(os.Stderr, "events: dsn\n")
		os.Exit(1)
	}
	pc.MaxConns = 8
	pool, err := pgxpool.NewWithConfig(ctx, pc)
	if err != nil {
		fmt.Fprintf(os.Stderr, "events: pool\n")
		os.Exit(1)
	}
	defer pool.Close()

	if err := db.RunMigrations(ctx, pool, db.ResolveMigrationsDir()); err != nil {
		fmt.Fprintf(os.Stderr, "events: migrate\n")
		os.Exit(1)
	}

	owner := fmt.Sprintf("events-%d", os.Getpid())
	disp := events.NewDispatcher(pool, cfg.Env, owner)
	mux := http.NewServeMux()
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
		writeJSON(w, http.StatusOK, map[string]any{"status": "ok", "env": cfg.Env, "owner": owner})
	})
	mux.HandleFunc("GET /ready", func(w http.ResponseWriter, r *http.Request) {
		if err := events.PingPool(r.Context(), pool); err != nil {
			writeJSON(w, http.StatusServiceUnavailable, map[string]any{"status": "not_ready", "code": "DB"})
			return
		}
		if err := events.SchemaReady(r.Context(), pool); err != nil {
			writeJSON(w, http.StatusServiceUnavailable, map[string]any{"status": "not_ready", "code": "SCHEMA"})
			return
		}
		reg := events.NewRegistry()
		if err := reg.Register(events.ConsumerNotify); err != nil {
			writeJSON(w, http.StatusServiceUnavailable, map[string]any{"status": "not_ready", "code": "REGISTRY"})
			return
		}
		writeJSON(w, http.StatusOK, map[string]any{"status": "ok"})
	})
	mux.HandleFunc("GET /live", func(w http.ResponseWriter, _ *http.Request) {
		if !disp.Alive(30 * time.Second) {
			writeJSON(w, http.StatusServiceUnavailable, map[string]any{"status": "stale"})
			return
		}
		writeJSON(w, http.StatusOK, map[string]any{"status": "ok"})
	})
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		m := disp.Metrics
		writeJSON(w, http.StatusOK, map[string]any{
			"claimed": m.Claimed.Load(), "acked": m.Acked.Load(), "failed": m.Failed.Load(),
			"fenced": m.Fenced.Load(), "dead": m.Dead.Load(), "quarantine": m.Quarantine.Load(),
			"heartbeat": m.Heartbeat.Load(),
		})
	})

	srv := &http.Server{Addr: fmt.Sprintf("%s:%d", cfg.HTTPHost, port), Handler: mux}
	go func() {
		consumers := []string{events.ConsumerNotify, events.ConsumerProjection, events.ConsumerWorkflowAdvance}
		_ = disp.Serve(ctx, consumers)
	}()
	go func() {
		<-ctx.Done()
		shut, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		disp.StopClaim()
		_ = srv.Shutdown(shut)
	}()
	fmt.Printf("events worker env=%s listen=%s:%d owner=%s mode=cluster-table\n", cfg.Env, cfg.HTTPHost, port, owner)
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		os.Exit(1)
	}
}

func parsePort(v string) (int, error) {
	var p int
	_, err := fmt.Sscanf(v, "%d", &p)
	if err != nil || p <= 0 {
		return 0, os.ErrInvalid
	}
	return p, nil
}

func writeJSON(w http.ResponseWriter, status int, body map[string]any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(body)
}