CHORUS/pkg/shutdown/components.go

package shutdown

import (
	"context"
	"fmt"
	"net/http"
	"time"
)

// HTTPServerComponent wraps an HTTP server for graceful shutdown
type HTTPServerComponent struct {
	name     string
	server   *http.Server
	priority int
}

// NewHTTPServerComponent creates a new HTTP server component
func NewHTTPServerComponent(name string, server *http.Server, priority int) *HTTPServerComponent {
	return &HTTPServerComponent{
		name:     name,
		server:   server,
		priority: priority,
	}
}

func (h *HTTPServerComponent) Name() string {
	return h.name
}

func (h *HTTPServerComponent) Priority() int {
	return h.priority
}

func (h *HTTPServerComponent) CanForceStop() bool {
	return true
}

func (h *HTTPServerComponent) Shutdown(ctx context.Context) error {
	if h.server == nil {
		return nil
	}
	
	return h.server.Shutdown(ctx)
}

// P2PNodeComponent wraps a P2P node for graceful shutdown
type P2PNodeComponent struct {
	name     string
	closer   func() error
	priority int
}

// NewP2PNodeComponent creates a new P2P node component
func NewP2PNodeComponent(name string, closer func() error, priority int) *P2PNodeComponent {
	return &P2PNodeComponent{
		name:     name,
		closer:   closer,
		priority: priority,
	}
}

func (p *P2PNodeComponent) Name() string {
	return p.name
}

func (p *P2PNodeComponent) Priority() int {
	return p.priority
}

func (p *P2PNodeComponent) CanForceStop() bool {
	return true
}

func (p *P2PNodeComponent) Shutdown(ctx context.Context) error {
	if p.closer == nil {
		return nil
	}
	
	// P2P nodes typically need time to disconnect gracefully
	done := make(chan error, 1)
	go func() {
		done <- p.closer()
	}()
	
	select {
	case err := <-done:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// DatabaseComponent wraps a database connection for graceful shutdown
type DatabaseComponent struct {
	name     string
	closer   func() error
	priority int
}

// NewDatabaseComponent creates a new database component
func NewDatabaseComponent(name string, closer func() error, priority int) *DatabaseComponent {
	return &DatabaseComponent{
		name:     name,
		closer:   closer,
		priority: priority,
	}
}

func (d *DatabaseComponent) Name() string {
	return d.name
}

func (d *DatabaseComponent) Priority() int {
	return d.priority
}

func (d *DatabaseComponent) CanForceStop() bool {
	return false // Databases shouldn't be force-stopped
}

func (d *DatabaseComponent) Shutdown(ctx context.Context) error {
	if d.closer == nil {
		return nil
	}
	
	return d.closer()
}

// ElectionManagerComponent wraps an election manager for graceful shutdown
type ElectionManagerComponent struct {
	name     string
	stopper  func()
	priority int
}

// NewElectionManagerComponent creates a new election manager component
func NewElectionManagerComponent(name string, stopper func(), priority int) *ElectionManagerComponent {
	return &ElectionManagerComponent{
		name:     name,
		stopper:  stopper,
		priority: priority,
	}
}

func (e *ElectionManagerComponent) Name() string {
	return e.name
}

func (e *ElectionManagerComponent) Priority() int {
	return e.priority
}

func (e *ElectionManagerComponent) CanForceStop() bool {
	return true
}

func (e *ElectionManagerComponent) Shutdown(ctx context.Context) error {
	if e.stopper == nil {
		return nil
	}
	
	// Election managers need special handling to transfer leadership
	done := make(chan struct{})
	go func() {
		e.stopper()
		close(done)
	}()
	
	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// PubSubComponent wraps a PubSub system for graceful shutdown
type PubSubComponent struct {
	name     string
	closer   func() error
	priority int
}

// NewPubSubComponent creates a new PubSub component
func NewPubSubComponent(name string, closer func() error, priority int) *PubSubComponent {
	return &PubSubComponent{
		name:     name,
		closer:   closer,
		priority: priority,
	}
}

func (p *PubSubComponent) Name() string {
	return p.name
}

func (p *PubSubComponent) Priority() int {
	return p.priority
}

func (p *PubSubComponent) CanForceStop() bool {
	return true
}

func (p *PubSubComponent) Shutdown(ctx context.Context) error {
	if p.closer == nil {
		return nil
	}
	
	return p.closer()
}

// MonitoringComponent wraps a monitoring system for graceful shutdown
type MonitoringComponent struct {
	name     string
	closer   func() error
	priority int
}

// NewMonitoringComponent creates a new monitoring component
func NewMonitoringComponent(name string, closer func() error, priority int) *MonitoringComponent {
	return &MonitoringComponent{
		name:     name,
		closer:   closer,
		priority: priority,
	}
}

func (m *MonitoringComponent) Name() string {
	return m.name
}

func (m *MonitoringComponent) Priority() int {
	return m.priority
}

func (m *MonitoringComponent) CanForceStop() bool {
	return true
}

func (m *MonitoringComponent) Shutdown(ctx context.Context) error {
	if m.closer == nil {
		return nil
	}
	
	return m.closer()
}

// GenericComponent provides a generic wrapper for any component with a close function
type GenericComponent struct {
	name         string
	closer       func() error
	priority     int
	canForceStop bool
	shutdownFunc func(ctx context.Context) error
}

// NewGenericComponent creates a new generic component
func NewGenericComponent(name string, priority int, canForceStop bool) *GenericComponent {
	return &GenericComponent{
		name:         name,
		priority:     priority,
		canForceStop: canForceStop,
	}
}

// SetCloser sets a simple closer function
func (g *GenericComponent) SetCloser(closer func() error) *GenericComponent {
	g.closer = closer
	return g
}

// SetShutdownFunc sets a context-aware shutdown function
func (g *GenericComponent) SetShutdownFunc(shutdownFunc func(ctx context.Context) error) *GenericComponent {
	g.shutdownFunc = shutdownFunc
	return g
}

func (g *GenericComponent) Name() string {
	return g.name
}

func (g *GenericComponent) Priority() int {
	return g.priority
}

func (g *GenericComponent) CanForceStop() bool {
	return g.canForceStop
}

func (g *GenericComponent) Shutdown(ctx context.Context) error {
	if g.shutdownFunc != nil {
		return g.shutdownFunc(ctx)
	}
	
	if g.closer != nil {
		// Wrap simple closer in context-aware function
		done := make(chan error, 1)
		go func() {
			done <- g.closer()
		}()
		
		select {
		case err := <-done:
			return err
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	
	return nil
}

// WorkerPoolComponent manages a pool of workers for graceful shutdown
type WorkerPoolComponent struct {
	name         string
	stopCh       chan struct{}
	workers      int
	priority     int
	shutdownTime time.Duration
}

// NewWorkerPoolComponent creates a new worker pool component
func NewWorkerPoolComponent(name string, stopCh chan struct{}, workers int, priority int) *WorkerPoolComponent {
	return &WorkerPoolComponent{
		name:         name,
		stopCh:       stopCh,
		workers:      workers,
		priority:     priority,
		shutdownTime: 10 * time.Second,
	}
}

func (w *WorkerPoolComponent) Name() string {
	return fmt.Sprintf("%s (workers: %d)", w.name, w.workers)
}

func (w *WorkerPoolComponent) Priority() int {
	return w.priority
}

func (w *WorkerPoolComponent) CanForceStop() bool {
	return true
}

func (w *WorkerPoolComponent) Shutdown(ctx context.Context) error {
	if w.stopCh == nil {
		return nil
	}
	
	// Signal workers to stop
	close(w.stopCh)
	
	// Wait for workers to finish with timeout
	timeout := w.shutdownTime
	if deadline, ok := ctx.Deadline(); ok {
		if remaining := time.Until(deadline); remaining < timeout {
			timeout = remaining
		}
	}
	
	// In a real implementation, you would wait for workers to signal completion
	select {
	case <-time.After(timeout):
		return fmt.Errorf("workers did not shut down within %v", timeout)
	case <-ctx.Done():
		return ctx.Err()
	}
}