A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
13 kB · 380 lines
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381package controller
import ( "context" "errors" "fmt" "log" "net/http" "os" "path/filepath" "strings" "sync" "time"
appnotify "forgejo.alexma.top/alexma233/composia/internal/app/notify" "forgejo.alexma.top/alexma233/composia/internal/core/config" corenotify "forgejo.alexma.top/alexma233/composia/internal/core/notify" "forgejo.alexma.top/alexma233/composia/internal/core/repo" "forgejo.alexma.top/alexma233/composia/internal/core/task" "forgejo.alexma.top/alexma233/composia/internal/platform/rpcutil" "forgejo.alexma.top/alexma233/composia/internal/platform/store")
func runControllerRuntime(ctx context.Context, cfg *config.ControllerConfig, reload func(context.Context) error) error { if err := os.MkdirAll(cfg.StateDir, 0o750); err != nil { return fmt.Errorf("create controller state_dir %q: %w", cfg.StateDir, err) } if err := os.MkdirAll(cfg.LogDir, 0o750); err != nil { return fmt.Errorf("create controller log_dir %q: %w", cfg.LogDir, err) } if err := os.MkdirAll(filepath.Join(cfg.LogDir, "tasks"), 0o750); err != nil { return fmt.Errorf("create controller task log_dir %q: %w", filepath.Join(cfg.LogDir, "tasks"), err) } if err := repo.ValidateWorkingTree(cfg.RepoDir); err != nil { return err } db, err := store.Open(cfg.StateDir) if err != nil { return err } runtimeCtx, cancelRuntime := context.WithCancel(ctx) var background sync.WaitGroup defer func() { cancelRuntime() background.Wait() _ = db.Close() }() startBackground := func(run func()) { background.Add(1) go func() { defer background.Done() run() }() }
nodeIDs := make([]string, 0, len(cfg.Nodes)) availableNodeIDs := make(map[string]struct{}, len(cfg.Nodes)) for _, node := range cfg.Nodes { nodeIDs = append(nodeIDs, node.ID) availableNodeIDs[node.ID] = struct{}{} } if err := db.SyncConfiguredNodes(runtimeCtx, nodeIDs); err != nil { return err } if err := db.MarkOfflineNodesBefore(runtimeCtx, time.Now().Add(-heartbeatOfflineAfter)); err != nil { return err } if _, _, err := db.SweepExpiredTaskExecutions(runtimeCtx, time.Now().UTC()); err != nil { return err } repoMu := &sync.Mutex{} taskQueue := newTaskQueueNotifier() taskResults := newTaskResultNotifier() dockerQueries := newDockerQueryBroker() execManager := newExecTunnelManager() logManager := newContainerLogTunnelManager() startBackground(func() { execManager.runSessionSweeper(runtimeCtx, execSessionSweepInterval) }) notifier, err := appnotify.New(cfg.Notifications) if err != nil { return fmt.Errorf("initialize notifications: %w", err) } reportServer := &agentReportServer{db: db, cfg: cfg, availableNodeIDs: availableNodeIDs, logState: &taskLogAckState{confirmedBy: make(map[string]uint64)}, taskQueue: taskQueue, taskResults: taskResults, dockerQueries: dockerQueries, execManager: execManager, logManager: logManager, repoMu: repoMu, notifier: notifier} startBackground(func() { runTaskOutbox(runtimeCtx, reportServer) })
services, err := repo.DiscoverServices(cfg.RepoDir, availableNodeIDs) if err != nil { return err } declaredServices := make(map[string][]string, len(services)) for _, service := range services { declaredServices[service.Name] = append([]string(nil), service.TargetNodes...) } if err := db.SyncDeclaredServices(runtimeCtx, declaredServices); err != nil { return err } startBackground(func() { runControllerTasks(runtimeCtx, &controllerTaskExecutor{db: db, cfg: cfg, availableNodeIDs: availableNodeIDs, taskQueue: taskQueue, dnsProviders: defaultDNSProviderFactory{}, taskResults: taskResults, repoMu: repoMu, notifier: notifier}) })
mux := http.NewServeMux() agentTokens := cfg.NodeTokenMap() accessTokens := cfg.EnabledAccessTokenMap() controllerStartedAt := time.Now().UTC()
agentInterceptor := rpcutil.NewServerBearerAuthInterceptor(func(token string) (string, error) { nodeID, ok := agentTokens[token] if !ok { return "", errors.New("invalid agent token") } return nodeID, nil })
accessInterceptor := rpcutil.NewServerBearerAuthInterceptor(func(token string) (string, error) { name, ok := accessTokens[token] if !ok { return "", errors.New("invalid access token") } return name, nil }) registerAgentHandlers(mux, cfg, db, agentInterceptor, taskQueue, taskResults, dockerQueries, execManager, logManager, repoMu, notifier) registerAccessHandlers(mux, cfg, db, accessInterceptor, availableNodeIDs, taskQueue, taskResults, dockerQueries, execManager, logManager, repoMu, reload, notifier) registerMetricsHandler(mux, db, accessTokens, controllerStartedAt) registerAlertmanagerHandler(mux, cfg.Notifications, notifier) mux.HandleFunc(rpcutil.ControllerExecWSPath, execManager.handleWebsocket)
protocols := new(http.Protocols) protocols.SetHTTP1(true) protocols.SetUnencryptedHTTP2(true)
server := &http.Server{ Addr: cfg.ListenAddr, Handler: mux, Protocols: protocols, ReadHeaderTimeout: 5 * time.Second, }
startBackground(func() { sweepOfflineNodes(runtimeCtx, db, notifier) }) startBackground(func() { sweepTaskExecutionLeases(runtimeCtx, db, taskQueue) }) startBackground(func() { autoPullRepo(runtimeCtx, cfg, db, availableNodeIDs, repoMu, taskQueue) }) startBackground(func() { runScheduledTasks(runtimeCtx, db, cfg, availableNodeIDs, taskQueue, repoMu) }) startBackground(func() { <-runtimeCtx.Done() shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if server.Shutdown(shutdownCtx) != nil { // Active streams must disconnect before the replacement runtime starts. _ = server.Close() } })
log.Printf("composia controller parsed %d declared services", len(services)) log.Printf("composia controller listening on %s", cfg.ListenAddr) err = server.ListenAndServe() cancelRuntime() if err != nil && !errors.Is(err, http.ErrServerClosed) { return fmt.Errorf("run controller server: %w", err) } return nil}
func sweepTaskExecutionLeases(ctx context.Context, db *store.DB, taskQueue *taskQueueNotifier) { ticker := time.NewTicker(offlineSweepInterval) defer ticker.Stop() for { select { case <-ctx.Done(): return case <-ticker.C: requeued, lost, err := db.SweepExpiredTaskExecutions(ctx, time.Now().UTC()) if err != nil { log.Printf("task execution sweep failed: %v", err) continue } if requeued > 0 { notifyTaskQueue(taskQueue) } if lost > 0 { log.Printf("task execution leases lost: count=%d", lost) } } }}
func sweepOfflineNodes(ctx context.Context, db *store.DB, notifier *appnotify.Notifier) { ticker := time.NewTicker(offlineSweepInterval) defer ticker.Stop()
for { select { case <-ctx.Done(): return case <-ticker.C: before, err := buildNodeSnapshotMap(ctx, db) if err != nil { log.Printf("offline sweep snapshot failed: %v", err) continue } if err := db.MarkOfflineNodesBefore(ctx, time.Now().Add(-heartbeatOfflineAfter)); err != nil { log.Printf("offline sweep failed: %v", err) continue } after, err := buildNodeSnapshotMap(ctx, db) if err != nil { log.Printf("offline sweep snapshot refresh failed: %v", err) continue } for nodeID, previous := range before { current, ok := after[nodeID] if !ok || !previous.IsOnline || current.IsOnline { continue } dispatchNodeNotification(notifier, corenotify.EventNodeOffline, current) } } }}
func autoPullRepo(ctx context.Context, cfg *config.ControllerConfig, db *store.DB, availableNodeIDs map[string]struct{}, repoMu *sync.Mutex, taskQueue *taskQueueNotifier) { if cfg.Git == nil || strings.TrimSpace(cfg.Git.RemoteURL) == "" || strings.TrimSpace(cfg.Git.PullInterval) == "" { return } interval, err := time.ParseDuration(strings.TrimSpace(cfg.Git.PullInterval)) if err != nil { log.Printf("auto-pull: invalid pull_interval %q: %v", cfg.Git.PullInterval, err) return } ticker := time.NewTicker(interval) defer ticker.Stop()
for { select { case <-ctx.Done(): return case <-ticker.C: repoMu.Lock() previousRevision, _ := repo.CurrentRevision(cfg.RepoDir) _, pullErr := autoPullFetchAndFastForward(ctx, cfg, db) repoMu.Unlock()
if pullErr != nil { log.Printf("auto-pull failed: %v", pullErr) continue } newRevision, _ := repo.CurrentRevision(cfg.RepoDir) if newRevision != previousRevision { log.Printf("auto-pull: repo updated from %s to %s", shortRevision(previousRevision), shortRevision(newRevision)) if err := refreshDeclaredServices(ctx, db, cfg, availableNodeIDs); err != nil { log.Printf("auto-pull: refresh declared services failed: %v", err) } handleAutoDeploy(ctx, cfg, db, availableNodeIDs, previousRevision, newRevision, taskQueue) } } }}
func handleAutoDeploy(ctx context.Context, cfg *config.ControllerConfig, db *store.DB, availableNodeIDs map[string]struct{}, oldRevision, newRevision string, taskQueue *taskQueueNotifier) { changedFiles, err := repo.DiffChangedFiles(cfg.RepoDir, oldRevision, newRevision) if err != nil { log.Printf("auto-deploy: diff changed files failed: %v", err) return } if len(changedFiles) == 0 { return } affectedNames, err := repo.AffectedServicesFromChangedFiles(cfg.RepoDir, changedFiles) if err != nil { log.Printf("auto-deploy: find affected services failed: %v", err) return } if len(affectedNames) == 0 { return } log.Printf("auto-deploy: detected %d affected services: %s", len(affectedNames), strings.Join(affectedNames, ", "))
for _, serviceName := range affectedNames { service, err := repo.FindService(cfg.RepoDir, availableNodeIDs, serviceName) if err != nil { log.Printf("auto-deploy: lookup service %q failed: %v", serviceName, err) continue } isInfra := service.Meta.IsInfra() shouldAutoDeploy := shouldAutoDeployService(cfg, service, isInfra)
if err := db.SetServicePendingDeploy(ctx, serviceName, service.TargetNodes, newRevision); err != nil { log.Printf("auto-deploy: set pending deploy for %q failed: %v", serviceName, err) } if shouldAutoDeploy { log.Printf("auto-deploy: triggering deploy for %q", serviceName) server := &serviceCommandServer{ db: db, cfg: cfg, availableNodeIDs: availableNodeIDs, taskQueue: taskQueue, } if _, err := server.createServiceTaskWithOptions(ctx, serviceName, service.TargetNodes, task.TypeDeploy, nil, serviceTaskCreateOptions{Source: task.SourceAutoDeploy}); err != nil { log.Printf("auto-deploy: create deploy task for %q failed: %v", serviceName, err) } } }}
func shouldAutoDeployService(cfg *config.ControllerConfig, service repo.Service, isInfra bool) bool { if !service.Meta.AutoDeployEnabled() { return false } if cfg.AutoDeploy == nil { return false } if isInfra { return cfg.AutoDeploy.Infra } return cfg.AutoDeploy.Services}
func autoPullFetchAndFastForward(ctx context.Context, cfg *config.ControllerConfig, db *store.DB) (store.RepoSyncState, error) { cleanWorktree, err := repo.IsCleanWorkingTree(cfg.RepoDir) if err != nil { return store.RepoSyncState{SyncStatus: store.RepoSyncStatusUnknown}, fmt.Errorf("check worktree clean: %w", err) } if !cleanWorktree { return store.RepoSyncState{SyncStatus: store.RepoSyncStatusLocalOnly}, nil } branch := strings.TrimSpace(cfg.Git.Branch) if branch == "" { branch, err = repo.CurrentBranch(cfg.RepoDir) if err != nil { return store.RepoSyncState{SyncStatus: store.RepoSyncStatusUnknown}, fmt.Errorf("get current branch: %w", err) } } authUsername := "" authToken := "" if cfg.Git.Auth != nil { authUsername = strings.TrimSpace(cfg.Git.Auth.Username) authToken = strings.TrimSpace(cfg.Git.Auth.Token) } previousState, err := db.GetRepoSyncState(ctx) if err != nil { previousState = store.RepoSyncState{} } pulledAt := time.Now().UTC().Format(time.RFC3339) if err := repo.FetchAndFastForward(cfg.RepoDir, strings.TrimSpace(cfg.Git.RemoteURL), branch, authUsername, authToken); err != nil { state := store.RepoSyncState{ SyncStatus: store.RepoSyncStatusPullFailed, LastSyncError: err.Error(), LastSuccessfulPullAt: previousState.LastSuccessfulPullAt, } if persistErr := db.UpsertRepoSyncState(ctx, state); persistErr != nil { return store.RepoSyncState{}, fmt.Errorf("persist pull failed state: %w", persistErr) } return state, fmt.Errorf("fetch and fast-forward: %w", err) } state := store.RepoSyncState{ SyncStatus: store.RepoSyncStatusSynced, LastSyncError: "", LastSuccessfulPullAt: pulledAt, } if err := db.UpsertRepoSyncState(ctx, state); err != nil { return store.RepoSyncState{}, fmt.Errorf("persist synced state: %w", err) } return state, nil}
func refreshDeclaredServices(ctx context.Context, db *store.DB, cfg *config.ControllerConfig, availableNodeIDs map[string]struct{}) error { services, err := repo.DiscoverServices(cfg.RepoDir, availableNodeIDs) if err != nil { return err } declaredServices := make(map[string][]string, len(services)) for _, service := range services { declaredServices[service.Name] = append([]string(nil), service.TargetNodes...) } return db.SyncDeclaredServices(ctx, declaredServices)}
func shortRevision(revision string) string { if len(revision) <= 8 { return revision } return revision[:8]}