backend/server/server.go

Rama main · réplica de lectura del repositorio original.

// Package server implements the API server for Bytebase.
package server

import (
	"context"
	"fmt"
	"log/slog"
	"net"
	"net/http"
	"path"
	"sync"
	"time"

	"github.com/labstack/echo/v5"
	"github.com/pkg/errors"

	directorysync "github.com/bytebase/bytebase/backend/api/directory-sync"
	"github.com/bytebase/bytebase/backend/api/lsp"
	"github.com/bytebase/bytebase/backend/api/mcp"
	"github.com/bytebase/bytebase/backend/api/oauth2"
	stripeapi "github.com/bytebase/bytebase/backend/api/stripe"
	"github.com/bytebase/bytebase/backend/common"
	"github.com/bytebase/bytebase/backend/common/log"
	"github.com/bytebase/bytebase/backend/component/bus"
	"github.com/bytebase/bytebase/backend/component/config"
	"github.com/bytebase/bytebase/backend/component/dbfactory"
	"github.com/bytebase/bytebase/backend/component/iam"
	"github.com/bytebase/bytebase/backend/component/productmetrics"
	"github.com/bytebase/bytebase/backend/component/review"
	"github.com/bytebase/bytebase/backend/component/sample"
	"github.com/bytebase/bytebase/backend/component/sample/saas"
	"github.com/bytebase/bytebase/backend/component/sample/selfhost"
	"github.com/bytebase/bytebase/backend/component/sheet"
	"github.com/bytebase/bytebase/backend/component/telemetry"
	"github.com/bytebase/bytebase/backend/component/webhook"
	"github.com/bytebase/bytebase/backend/enterprise"
	storepb "github.com/bytebase/bytebase/backend/generated-go/store"
	v1pb "github.com/bytebase/bytebase/backend/generated-go/v1"
	"github.com/bytebase/bytebase/backend/migrator"
	"github.com/bytebase/bytebase/backend/resources/postgres"
	"github.com/bytebase/bytebase/backend/runner/cleaner"
	"github.com/bytebase/bytebase/backend/runner/heartbeat"
	"github.com/bytebase/bytebase/backend/runner/monitor"
	"github.com/bytebase/bytebase/backend/runner/notifylistener"
	"github.com/bytebase/bytebase/backend/runner/plancheck"
	"github.com/bytebase/bytebase/backend/runner/reviewrun"
	samplerunner "github.com/bytebase/bytebase/backend/runner/sample"
	"github.com/bytebase/bytebase/backend/runner/schemasync"
	"github.com/bytebase/bytebase/backend/runner/taskrun"
	"github.com/bytebase/bytebase/backend/store"
)

const (
	// webhookAPIPrefix is the API prefix for Bytebase webhook.
	webhookAPIPrefix = "/hook"
	scimAPIPrefix    = "/scim"
	// lspAPI is the API for Bytebase Language Server Protocol.
	lspAPI                 = "/lsp"
	gracefulShutdownPeriod = 10 * time.Second
	// replicaDeregisterTimeout bounds the heartbeat delete that runs after the
	// graceful shutdown deadline may already have passed.
	replicaDeregisterTimeout = 3 * time.Second
)

// Server is the Bytebase server.
type Server struct {
	// Asynchronous runners.
	taskScheduler      *taskrun.Scheduler
	planCheckScheduler *plancheck.Scheduler
	reviewRunScheduler *reviewrun.Scheduler
	schemaSyncer       *schemasync.Syncer
	approvalRunner     *review.Runner
	notifyListener     *notifylistener.Listener
	dataCleaner        *cleaner.DataCleaner
	heartbeatRunner    *heartbeat.Runner
	sampleRunner       *samplerunner.Runner
	runnerWG           sync.WaitGroup

	webhookManager *webhook.Manager
	iamManager     *iam.Manager
	sampleManager  sample.Manager

	licenseService *enterprise.LicenseService

	profile    *config.Profile
	echoServer *echo.Echo
	httpServer *http.Server
	lspServer  *lsp.Server
	store      *store.Store
	dbFactory  *dbfactory.DBFactory
	startedTS  int64

	// PG server stoppers.
	stopper []func()

	// bus is the message bus for inter-component communication within the server.
	bus *bus.Bus

	// boot specifies that whether the server boot correctly
	cancel context.CancelFunc
}

// NewServer creates a server.
func NewServer(ctx context.Context, profile *config.Profile) (*Server, error) {
	s := &Server{
		profile:   profile,
		startedTS: time.Now().Unix(),
	}

	// Display config
	slog.Info("-----Config BEGIN-----")
	slog.Info(fmt.Sprintf("mode=%s", profile.Mode))
	slog.Info(fmt.Sprintf("dataDir=%s", profile.DataDir))
	slog.Info(fmt.Sprintf("replicaID=%s", profile.ReplicaID))
	slog.Info("-----Config END-------")

	serverStarted := false
	defer func() {
		if !serverStarted {
			_ = s.Shutdown(ctx)
		}
	}()

	if profile.HA && profile.UseEmbedDB() {
		return nil, errors.New("HA mode requires external PostgreSQL (set PG_URL environment variable)")
	}

	var pgURL string
	if profile.UseEmbedDB() {
		pgDataDir := path.Join(profile.DataDir, "pgdata")

		stopper, err := postgres.StartMetadataInstance(ctx, pgDataDir, profile.DatastorePort, profile.Mode)
		if err != nil {
			return nil, err
		}
		s.stopper = append(s.stopper, stopper)
		pgURL = fmt.Sprintf("host=%s port=%d user=bb database=bb", postgres.GetPostgresSocketDir(), profile.DatastorePort)
	} else {
		pgURL = profile.PgURL
	}

	// Connect to the instance that stores bytebase's own metadata.
	stores, err := store.New(ctx, pgURL, !profile.HA)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to new store")
	}

	if err := migrator.MigrateSchema(ctx, stores.GetDB()); err != nil {
		stores.Close()
		return nil, errors.Wrapf(err, "failed to migrate schema")
	}
	s.store = stores
	sheetManager := sheet.NewManager()

	s.licenseService, err = enterprise.NewLicenseService(profile.Mode, stores, profile.SaaS, profile.LicensePrivateKey)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create license service")
	}

	logSetup := &storepb.WorkspaceProfileSetting{
		EnableAuditLogStdout:   false,
		EnableMetricCollection: true,
		EnableDebug:            profile.Debug,
	}
	var workspaceID string
	if !s.profile.SaaS {
		// Load workspace-dependent settings if workspace exists.
		// On first boot (no workspace yet), these remain at defaults and get
		// initialized when the workspace is created and settings are updated via API.
		if workspaceID, _ = stores.GetWorkspaceID(ctx); workspaceID != "" {
			if workspaceProfileSetting, err := s.store.GetWorkspaceProfileSetting(ctx, workspaceID); err == nil {
				logSetup = workspaceProfileSetting
				if logSetup.GetEnableAuditLogStdout() && s.licenseService.IsFeatureEnabled(ctx, workspaceID, v1pb.PlanFeature_FEATURE_AUDIT_LOG) == nil {
					s.profile.RuntimeEnableAuditLogStdout.Store(true)
				}
			}
		}
	}

	profile.RuntimeDebug.Store(logSetup.EnableDebug)
	if logSetup.EnableDebug {
		log.LogLevel.Set(slog.LevelDebug)
	}
	telemetry.InitGlobalReporter(
		profile.Version,
		profile.GitCommit,
		profile.Mode,
		logSetup.GetEnableMetricCollection(),
	)
	if !s.profile.SaaS && profile.Mode == common.ReleaseModeProd && logSetup.GetEnableMetricCollection() && workspaceID != "" {
		go func() {
			reportSQLReviewConfigSnapshot(context.WithoutCancel(ctx), stores, workspaceID)
		}()
	}

	s.bus, err = bus.New()
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create message bus")
	}

	// The auth secret is a global infrastructure value stored in server_config.
	secret, err := s.store.GetAuthSecret(ctx)
	if err != nil {
		return nil, errors.Wrap(err, "failed to get auth secret")
	}

	s.iamManager, err = iam.NewManager(stores, s.licenseService, profile.SaaS)
	if err := s.iamManager.ReloadCache(ctx); err != nil {
		return nil, errors.Wrapf(err, "failed to reload iam cache")
	}
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create iam manager")
	}
	s.webhookManager = webhook.NewManager(stores, profile)
	s.dbFactory = dbfactory.New(s.store, s.licenseService)

	// Configure echo server.
	s.echoServer = echo.New()

	productMetrics := productmetrics.New(stores, s.licenseService)
	s.schemaSyncer = schemasync.NewSyncer(stores, s.dbFactory, s.licenseService, productMetrics)
	if profile.SaaS {
		s.sampleManager = configureSampleManager(
			ctx,
			profile.SampleProjectInstancePgURL,
			stores,
			s.schemaSyncer,
			profile.ReplicaID,
		)
	} else {
		s.sampleManager = selfhost.NewManager(stores, profile, s.schemaSyncer, sample.ManagerOptions{ReplicaID: profile.ReplicaID})
	}
	if s.sampleManager != nil {
		s.sampleRunner = samplerunner.NewRunner(s.sampleManager)
		if workspaceID != "" {
			if err := s.sampleManager.Start(ctx, workspaceID); err != nil {
				slog.Warn("failed to start sample instances", log.BBError(err))
			}
		}
	}
	s.approvalRunner = review.NewRunner(stores, s.bus, s.webhookManager, s.licenseService)

	s.taskScheduler = taskrun.NewScheduler(stores, s.bus, s.webhookManager, s.licenseService, profile, productMetrics)
	s.taskScheduler.Register(storepb.Task_DATABASE_CREATE, taskrun.NewDatabaseCreateExecutor(stores, s.dbFactory, s.schemaSyncer))
	s.taskScheduler.Register(storepb.Task_DATABASE_MIGRATE, taskrun.NewDatabaseMigrateExecutor(stores, s.dbFactory, s.bus, s.schemaSyncer, profile))

	var ruleReviewCovers func(storepb.Engine) bool
	if profile.SQLReviewV2 {
		ruleReviewCovers = reviewrun.ReviewsEngine
	}
	combinedExecutor := plancheck.NewCombinedExecutor(stores, sheetManager, s.dbFactory, ruleReviewCovers)
	s.planCheckScheduler = plancheck.NewScheduler(stores, s.bus, combinedExecutor, s.licenseService, productMetrics)

	s.reviewRunScheduler = reviewrun.NewScheduler(stores, s.bus, profile, s.licenseService, productMetrics)
	s.reviewRunScheduler.Register(store.ReviewRunTypeRule, reviewrun.NewRuleExecutor(stores))
	s.reviewRunScheduler.Register(store.ReviewRunTypeAI, reviewrun.NewAIReviewExecutor(stores, ruleReviewCovers))
	s.notifyListener = notifylistener.NewListener(stores.GetDB(), s.bus)

	// Data cleaner
	s.dataCleaner = cleaner.NewDataCleaner(stores, s.licenseService)

	// Heartbeat runner
	s.heartbeatRunner = heartbeat.NewRunner(stores, profile)

	// LSP server.
	s.lspServer = lsp.NewServer(s.store, profile, secret, s.bus, s.iamManager, s.licenseService)

	directorySyncServer := directorysync.NewService(s.store, s.licenseService, s.iamManager, profile)
	oauth2Service := oauth2.NewService(stores, profile, secret)

	stripeWebhookHandler := stripeapi.NewWebhookHandler(s.store, s.licenseService, profile.StripeWebhookSecret)

	internalMCPHandler, err := configureGrpcRouters(ctx, s.echoServer, s.store, sheetManager, s.dbFactory, s.licenseService, s.profile, s.bus, s.schemaSyncer, s.webhookManager, s.iamManager, secret, s.sampleManager)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to configure gRPC routers")
	}
	mcpServer, err := mcp.NewServer(stores, profile, secret, internalMCPHandler, s.licenseService)
	if err != nil {
		return nil, errors.Wrapf(err, "failed to create MCP server")
	}
	configureEchoRouters(s.echoServer, s.lspServer, directorySyncServer, oauth2Service, mcpServer, stripeWebhookHandler, productMetrics, profile)

	serverStarted = true
	return s, nil
}

func configureSampleManager(
	ctx context.Context,
	targetURL string,
	stores *store.Store,
	syncer *schemasync.Syncer,
	replicaID string,
) sample.Manager {
	if targetURL == "" {
		return nil
	}
	manager, err := saas.NewManager(
		stores,
		targetURL,
		syncer,
		sample.ManagerOptions{ReplicaID: replicaID},
	)
	if err != nil {
		slog.Warn("invalid SAMPLE_PROJECT_INSTANCE_PG_URL; Sample Project Instance is disabled", log.BBError(err))
		return nil
	}
	if err := manager.CheckAvailable(ctx); err != nil {
		slog.Warn("Sample Project Instance target is temporarily unavailable", log.BBError(err))
	}
	return manager
}

// Run will run the server.
func (s *Server) Run(ctx context.Context, port int) error {
	ctx, cancel := context.WithCancel(ctx)
	s.cancel = cancel
	// runnerWG waits for all goroutines to complete.
	s.runnerWG.Add(1)
	go s.taskScheduler.Run(ctx, &s.runnerWG)
	s.runnerWG.Add(1)
	go s.schemaSyncer.Run(ctx, &s.runnerWG)
	s.runnerWG.Add(1)
	go s.approvalRunner.Run(ctx, &s.runnerWG)

	s.runnerWG.Add(1)
	go s.planCheckScheduler.Run(ctx, &s.runnerWG)

	s.runnerWG.Add(1)
	go s.reviewRunScheduler.Run(ctx, &s.runnerWG)

	s.runnerWG.Add(1)
	go s.dataCleaner.Run(ctx, &s.runnerWG)

	s.runnerWG.Add(1)
	go s.heartbeatRunner.Run(ctx, &s.runnerWG)

	if s.sampleRunner != nil {
		s.runnerWG.Add(1)
		go s.sampleRunner.Run(ctx, &s.runnerWG)
	}

	s.runnerWG.Add(1)
	go s.notifyListener.Run(ctx, &s.runnerWG)

	s.runnerWG.Add(1)
	mmm := monitor.NewMemoryMonitor(s.profile)
	go mmm.Run(ctx, &s.runnerWG)

	address := fmt.Sprintf(":%d", port)
	listener, err := net.Listen("tcp", address)
	if err != nil {
		return err
	}

	// Create HTTP server with H2C support (Echo v5)
	protocols := new(http.Protocols)
	protocols.SetHTTP1(true)
	protocols.SetUnencryptedHTTP2(true)
	s.httpServer = &http.Server{
		Addr:      address,
		Handler:   s.echoServer,
		Protocols: protocols,
	}

	go func() {
		if err := s.httpServer.Serve(listener); err != nil {
			if !errors.Is(err, http.ErrServerClosed) {
				slog.Error("http server listen error", log.BBError(err))
			}
		}
	}()

	return nil
}

// Shutdown will shut down the server.
func waitRunners(ctx context.Context, wg *sync.WaitGroup) error {
	stopped := make(chan struct{})
	go func() {
		wg.Wait()
		close(stopped)
	}()
	select {
	case <-stopped:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (s *Server) Shutdown(ctx context.Context) error {
	slog.Info("Stopping Bytebase...")
	slog.Info("Stopping web server...")

	ctx, cancel := context.WithTimeout(ctx, gracefulShutdownPeriod)
	defer cancel()

	// Cancel the worker
	if s.cancel != nil {
		s.cancel()
	}

	// Shutdown HTTP server
	if s.httpServer != nil {
		if err := s.httpServer.Shutdown(ctx); err != nil {
			slog.Error("Failed to shutdown HTTP server", log.BBError(err))
		}
	}

	// Wait for all runners to exit, but not past the shutdown deadline: one
	// blocked on a stopped consumer must not hold the process open. A runner
	// that outlives this gets an error from the closed store, not a nil handle.
	if err := waitRunners(ctx, &s.runnerWG); err != nil {
		slog.Warn("Runners did not stop before the shutdown deadline", log.BBError(err))
	}

	if s.store != nil {
		// Deregister only after runners stop: the heartbeat row is also the lease
		// peers check before failing this replica's task and review runs.
		deregisterCtx, cancelDeregister := context.WithTimeout(context.WithoutCancel(ctx), replicaDeregisterTimeout)
		if err := s.store.DeleteReplicaHeartbeat(deregisterCtx, s.profile.ReplicaID); err != nil {
			slog.Warn("Failed to delete replica heartbeat", slog.String("replicaID", s.profile.ReplicaID), log.BBError(err))
		}
		cancelDeregister()
	}

	// Close db connection
	if s.store != nil {
		if err := s.store.Close(); err != nil {
			return err
		}
	}

	// Shutdown sample instances
	if s.sampleManager != nil {
		s.sampleManager.Stop()
	}

	// Shutdown postgres instances.
	for _, stopper := range s.stopper {
		stopper()
	}

	return nil
}