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/app/eventsconsumer"
	"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)
	}

	h, err := eventsconsumer.Start(ctx, pool, cfg)
	if err != nil {
		fmt.Fprintf(os.Stderr, "events: start: %v\n", err)
		os.Exit(1)
	}
	defer h.Stop()

	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": h.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 !h.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 := h.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() {
		<-ctx.Done()
		shut, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()
		h.Stop()
		_ = srv.Shutdown(shut)
	}()
	fmt.Printf("events worker env=%s listen=%s:%d owner=%s mode=cluster-table\n", cfg.Env, cfg.HTTPHost, port, h.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)
}