pgsql-broker/pkg/broker/broker.go

package broker

import (
	"context"
	"fmt"
	"sync"
	"time"

	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/adapter"
	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/config"
	"git.warky.dev/wdevs/pgsql-broker/pkg/broker/metrics"
)

// Broker manages multiple database instances
type Broker struct {
	config    *config.Config
	logger    adapter.Logger
	version   string
	instances []*DatabaseInstance
	ctx       context.Context
	cancel    context.CancelFunc
	metrics   *metrics.Metrics
	server    *metrics.Server
	shutdown  bool
	mu        sync.RWMutex
}

// New creates a new broker that manages multiple database connections
func New(cfg *config.Config, logger adapter.Logger, version string) (*Broker, error) {
	ctx, cancel := context.WithCancel(context.Background())

	broker := &Broker{
		config:    cfg,
		logger:    logger.With("component", "broker"),
		version:   version,
		instances: make([]*DatabaseInstance, 0, len(cfg.Databases)),
		ctx:       ctx,
		cancel:    cancel,
	}
	if cfg.Broker.MetricsEnabled {
		broker.metrics = metrics.New()
		broker.metrics.SetDatabaseCount(len(cfg.Databases))
		addr := fmt.Sprintf("%s:%d", cfg.Broker.MetricsHost, cfg.Broker.MetricsPort)
		server, err := metrics.NewServer(broker.metrics, addr, broker.logger)
		if err != nil {
			return nil, err
		}
		broker.server = server
	}

	return broker, nil
}

// Start begins all database instances
func (b *Broker) Start() error {
	b.logger.Info("starting broker", "database_count", len(b.config.Databases))

	// Create and start an instance for each enabled database
	for i := range b.config.Databases {
		dbCfg := &b.config.Databases[i]
		if dbCfg.Disabled {
			b.logger.Info("skipping disabled database", "name", dbCfg.Name)
			continue
		}
		b.logger.Info("starting database instance", "name", dbCfg.Name, "host", dbCfg.Host, "database", dbCfg.Database)

		// Create database adapter
		dbAdapter := adapter.NewPostgresAdapter(dbCfg.ToPostgresConfig(), b.logger)

		// Create database instance
		instance, err := NewDatabaseInstance(b.config, dbCfg, dbAdapter, b.logger, b.version, b.ctx, b.metrics)
		if err != nil {
			// Stop any already-started instances
			b.stopInstances()
			return fmt.Errorf("failed to create database instance %d (%s): %w", i, dbCfg.Name, err)
		}

		// Start the instance
		if err := instance.Start(); err != nil {
			// Stop any already-started instances
			b.stopInstances()
			return fmt.Errorf("failed to start database instance %d (%s): %w", i, dbCfg.Name, err)
		}

		b.instances = append(b.instances, instance)
		b.logger.Info("database instance started", "name", dbCfg.Name, "instance_id", instance.ID)
	}

	if len(b.instances) == 0 {
		return fmt.Errorf("no enabled databases configured (all %d database(s) are disabled)", len(b.config.Databases))
	}

	b.logger.Info("broker started successfully", "database_instances", len(b.instances))
	if b.server != nil {
		b.server.Start()
	}
	return nil
}

// Stop gracefully stops all database instances
func (b *Broker) Stop() error {
	b.mu.Lock()
	if b.shutdown {
		b.mu.Unlock()
		return nil
	}
	b.shutdown = true
	b.mu.Unlock()

	b.logger.Info("stopping broker")

	// Cancel context
	b.cancel()

	// Stop all instances
	b.stopInstances()
	if b.server != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := b.server.Stop(ctx); err != nil {
			b.logger.Error("failed to stop metrics server", "error", err)
		}
	}

	b.logger.Info("broker stopped")
	return nil
}

// stopInstances stops all database instances
func (b *Broker) stopInstances() {
	var wg sync.WaitGroup
	for _, instance := range b.instances {
		wg.Add(1)
		go func(inst *DatabaseInstance) {
			defer wg.Done()
			defer adapter.RecoverAndLog(b.logger, "stop-instance-"+inst.DatabaseName)
			if err := inst.Stop(); err != nil {
				b.logger.Error("failed to stop instance", "name", inst.DatabaseName, "error", err)
			}
		}(instance)
	}
	wg.Wait()
}

// GetStats returns statistics for all database instances
func (b *Broker) GetStats() map[string]interface{} {
	b.mu.RLock()
	defer b.mu.RUnlock()

	stats := map[string]interface{}{
		"database_count": len(b.instances),
	}

	instanceStats := make(map[string]interface{})
	for _, instance := range b.instances {
		instanceStats[instance.DatabaseName] = instance.GetStats()
	}
	stats["instances"] = instanceStats

	return stats
}