cmd/example/main.go - path - kipitix/gracedown

cmd/example/main.go
251 строка7 KB
kipitix
upgrade
24 апр 2026, 00:54
24 апр 2026, 00:543842512
package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"
	"time"

	"github.com/kipitix/gracedown"
)

// CustomLogger implements the Logger interface
type CustomLogger struct {
	prefix string
}

func (l CustomLogger) Println(args ...any) {
	allArgs := append([]any{l.prefix}, args...)
	fmt.Println(allArgs...)
}

// Database represents a database connection
type Database struct {
	connected bool
}

func NewDatabase() *Database {
	return &Database{connected: false}
}

func (db *Database) Connect(ctx context.Context) error {
	fmt.Println("🔄 Connecting to PostgreSQL...")

	select {
	case <-ctx.Done():
		return fmt.Errorf("connection timeout: %w", ctx.Err())
	case <-time.After(2 * time.Second):
		// Uncomment to test EmergencyShutdown:
		// return errors.New("failed to connect to PostgreSQL: connection refused")
		db.connected = true
		fmt.Println("✅ PostgreSQL connected")
		return nil
	}
}

func (db *Database) Close(ctx context.Context) error {
	if !db.connected {
		return nil
	}
	fmt.Println("🔄 Closing PostgreSQL connections...")
	time.Sleep(1 * time.Second)
	db.connected = false
	fmt.Println("✅ PostgreSQL closed")
	return nil
}

// RedisCache represents a Redis connection
type RedisCache struct {
	connected bool
}

func NewRedisCache() *RedisCache {
	return &RedisCache{connected: false}
}

func (r *RedisCache) Connect(ctx context.Context) error {
	fmt.Println("🔄 Connecting to Redis...")

	select {
	case <-ctx.Done():
		return fmt.Errorf("connection timeout: %w", ctx.Err())
	case <-time.After(1 * time.Second):
		// Uncomment to test EmergencyShutdown:
		// return errors.New("failed to connect to Redis: no reachable nodes")
		r.connected = true
		fmt.Println("✅ Redis connected")
		return nil
	}
}

func (r *RedisCache) Close(ctx context.Context) error {
	if !r.connected {
		return nil
	}
	fmt.Println("🔄 Flushing Redis and closing connection...")
	time.Sleep(500 * time.Millisecond)
	r.connected = false
	fmt.Println("✅ Redis closed")
	return nil
}

// MessageQueue represents a RabbitMQ/Kafka connection
type MessageQueue struct {
	connected bool
}

func NewMessageQueue() *MessageQueue {
	return &MessageQueue{connected: false}
}

func (mq *MessageQueue) Connect(ctx context.Context) error {
	fmt.Println("🔄 Connecting to RabbitMQ...")

	select {
	case <-ctx.Done():
		return fmt.Errorf("connection timeout: %w", ctx.Err())
	case <-time.After(3 * time.Second):
		// Uncomment to test EmergencyShutdown:
		// return errors.New("failed to connect to RabbitMQ: authentication failed")
		mq.connected = true
		fmt.Println("✅ RabbitMQ connected")
		return nil
	}
}

func (mq *MessageQueue) Close(ctx context.Context) error {
	if !mq.connected {
		return nil
	}
	fmt.Println("🔄 Closing RabbitMQ channels and connection...")
	time.Sleep(1 * time.Second)
	mq.connected = false
	fmt.Println("✅ RabbitMQ closed")
	return nil
}

func main() {
	manager := gracedown.NewManager(
		gracedown.WithLogger(CustomLogger{prefix: "[GRACEDOWN]"}),
		gracedown.WithGlobalTimeout(45*time.Second),
	)

	// ========== INFRASTRUCTURE INITIALIZATION ==========
	initCtx, initCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer initCancel()

	db := NewDatabase()
	if err := db.Connect(initCtx); err != nil {
		fmt.Printf("❌ Failed to initialize database: %v\n", err)
		emergencyExit(manager, gracedown.ExitUnavailable)
	}
	manager.RegisterInfrastructure("PostgreSQL", 10*time.Second, db.Close)

	redis := NewRedisCache()
	if err := redis.Connect(initCtx); err != nil {
		fmt.Printf("❌ Failed to initialize Redis: %v\n", err)
		emergencyExit(manager, gracedown.ExitUnavailable)
	}
	manager.RegisterInfrastructure("Redis", 5*time.Second, redis.Close)

	mq := NewMessageQueue()
	if err := mq.Connect(initCtx); err != nil {
		fmt.Printf("❌ Failed to initialize Message Queue: %v\n", err)
		emergencyExit(manager, gracedown.ExitUnavailable)
	}
	manager.RegisterInfrastructure("RabbitMQ", 8*time.Second, mq.Close)

	// ========== INTERFACE LAYER ==========
	apiMux := http.NewServeMux()
	apiMux.HandleFunc("/api/users", func(w http.ResponseWriter, r *http.Request) {
		time.Sleep(5 * time.Second) // simulate slow request
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"users":[{"id":1,"name":"Alice"},{"id":2,"name":"Bob"}]}`))
	})
	apiMux.HandleFunc("/api/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(`{"status":"healthy","timestamp":"` + time.Now().Format(time.RFC3339) + `"}`))
	})

	apiServer := &http.Server{
		Addr:         ":8080",
		Handler:      apiMux,
		ReadTimeout:  10 * time.Second,
		WriteTimeout: 10 * time.Second,
	}

	manager.RegisterInterface("API HTTP Server", 15*time.Second, func(ctx context.Context) error {
		return apiServer.Shutdown(ctx)
	})

	metricsMux := http.NewServeMux()
	metricsMux.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte(`# HELP go_goroutines Number of goroutines
# TYPE go_goroutines gauge
go_goroutines 42
# HELP http_requests_total Total HTTP requests
# TYPE http_requests_total counter
http_requests_total{method="GET",endpoint="/api/users"} 1337
`))
	})
	metricsMux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"status":"ok","services":{"database":"healthy","redis":"healthy","rabbitmq":"healthy"}}`))
	})

	metricsServer := &http.Server{
		Addr:    ":9090",
		Handler: metricsMux,
	}

	manager.RegisterInterface("Metrics HTTP Server", 5*time.Second, func(ctx context.Context) error {
		return metricsServer.Shutdown(ctx)
	})

	// ========== START SERVERS ==========
	serverErrors := make(chan error, 2)

	go func() {
		fmt.Println("🚀 API Server starting on :8080")
		if err := apiServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			serverErrors <- fmt.Errorf("API Server error: %w", err)
		}
	}()

	go func() {
		fmt.Println("📊 Metrics Server starting on :9090")
		if err := metricsServer.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
			serverErrors <- fmt.Errorf("Metrics Server error: %w", err)
		}
	}()

	go func() {
		err := <-serverErrors
		fmt.Printf("❌ Server error: %v\n", err)
		emergencyExit(manager, gracedown.ExitIOErr)
	}()

	time.Sleep(100 * time.Millisecond)

	fmt.Println("✅ Application started successfully")
	fmt.Println("📝 Press Ctrl+C to initiate graceful shutdown")
	fmt.Println("📝 Or send SIGTERM signal: kill -TERM", os.Getpid())

	// Blocks until SIGINT or SIGTERM, then shuts down gracefully.
	if err := manager.WaitForSignalAndShutdown(); err != nil {
		fmt.Printf("❌ Shutdown completed with errors: %v\n", err)
		os.Exit(gracedown.ExitFailure)
	}
}

// emergencyExit shuts down all registered components and exits with the given code.
func emergencyExit(manager *gracedown.Manager, exitCode int) {
	if shutdownErr := manager.EmergencyShutdown(); shutdownErr != nil {
		fmt.Printf("❌ Emergency shutdown error: %v\n", shutdownErr)
	}
	os.Exit(exitCode)
}