mirror of
https://github.com/Buriburizaem0n/nezha_domains.git
synced 2026-09-19 17:50:12 +00:00
252 lines
9.5 KiB
Go
252 lines
9.5 KiB
Go
//go:build linux
|
|
|
|
package scenario
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"golang.org/x/sync/errgroup"
|
|
|
|
"github.com/nezhahq/nezha/integration/agentcompat/internal/client"
|
|
"github.com/nezhahq/nezha/integration/agentcompat/internal/dashboard"
|
|
)
|
|
|
|
type heldSessionSet struct {
|
|
mu sync.Mutex
|
|
plans []StressSessionPlan
|
|
sessions []heldSession
|
|
state heldSessionSetStateObserver
|
|
dependencies HeldSessionSetDependencies
|
|
baseline client.IOStreamState
|
|
healthContext context.Context
|
|
healthMu sync.Mutex
|
|
healthError error
|
|
healthErrors []error
|
|
healthDoneOnce sync.Once
|
|
healthDone chan struct{}
|
|
healthStop chan struct{}
|
|
healthStopOnce sync.Once
|
|
healthShutdown chan struct{}
|
|
healthShutdownOnce sync.Once
|
|
healthShutdownRequests chan heldHealthShutdownRequest
|
|
healthWG sync.WaitGroup
|
|
healthCoordinatorWG sync.WaitGroup
|
|
healthCoordinatorDoneOnce sync.Once
|
|
healthEvents chan heldHealthMessage
|
|
healthSnapshots []chan heldHealthSnapshotRequest
|
|
coordinatorDone chan struct{}
|
|
healthSnapshotRequestHook func(int)
|
|
healthSnapshotReplyHook func(int)
|
|
healthSnapshotOverrideHook func(int, heldHealthSnapshotRequest) *heldHealthSnapshot
|
|
healthSnapshotSendHook func(int)
|
|
healthEventHook func(heldHealthMessage)
|
|
healthClosureObservedHook func(int)
|
|
healthSnapshotAcceptedHook func(heldHealthSnapshot)
|
|
healthShutdownAcceptedHook func()
|
|
healthShutdownAcknowledgedHook func()
|
|
healthWatcherDone []chan struct{}
|
|
closing bool
|
|
closeOnce sync.Once
|
|
closeDone chan struct{}
|
|
closeError error
|
|
}
|
|
|
|
func NewHeldSessionSet(ctx context.Context, input HeldSessionSetInput) (*heldSessionSet, error) {
|
|
if ctx == nil {
|
|
return nil, ErrInvalidHeldSessionSetTopology
|
|
}
|
|
plans, err := validateHeldSessionSetPlans(input.Plan)
|
|
if err != nil {
|
|
return nil, redactHeldSessionSetError(err)
|
|
}
|
|
topology, err := validateHeldSessionSetTopology(input, plans)
|
|
if err != nil {
|
|
return nil, redactHeldSessionSetError(err)
|
|
}
|
|
dependencies := input.Dependencies
|
|
if (dependencies.Terminal == nil) != (dependencies.NAT == nil) || (dependencies.Terminal == nil) != (dependencies.FM == nil) || (dependencies.Terminal == nil) != (dependencies.Snapshot == nil) || (dependencies.Terminal == nil) != (dependencies.WaitState == nil) || (dependencies.Terminal == nil) != (dependencies.InspectAgent == nil) || (dependencies.Terminal == nil) != (dependencies.ObserveState == nil) {
|
|
return nil, redactHeldSessionSetError(ErrInvalidHeldSessionSetTopology)
|
|
}
|
|
if dependencies.Terminal == nil {
|
|
dependencies = defaultHeldSessionSetDependencies()
|
|
}
|
|
baseline, err := dependencies.Snapshot(ctx, topology.stateClient)
|
|
if err != nil {
|
|
return nil, redactHeldSessionSetError(err)
|
|
}
|
|
set := newHeldSessionSet(plans, topology.stateClient, baseline, dependencies, ctx)
|
|
set.healthSnapshotRequestHook = input.testHealthSnapshotRequestHook
|
|
set.healthSnapshotReplyHook = input.testHealthSnapshotReplyHook
|
|
set.healthSnapshotOverrideHook = input.testHealthSnapshotOverrideHook
|
|
set.healthSnapshotSendHook = input.testHealthSnapshotSendHook
|
|
set.healthEventHook = input.testHealthEventHook
|
|
set.healthClosureObservedHook = input.testHealthClosureObservedHook
|
|
set.healthSnapshotAcceptedHook = input.testHealthSnapshotAcceptedHook
|
|
set.healthShutdownAcceptedHook = input.testHealthShutdownAcceptedHook
|
|
set.healthShutdownAcknowledgedHook = input.testHealthShutdownAcknowledgedHook
|
|
if err := set.construct(ctx, topology); err != nil {
|
|
return nil, redactHeldSessionSetError(err)
|
|
}
|
|
if err := set.waitLive(ctx); err != nil {
|
|
return nil, redactHeldSessionSetError(errors.Join(err, set.Close(context.WithoutCancel(ctx))))
|
|
}
|
|
set.startHealthWatchers()
|
|
return set, nil
|
|
}
|
|
|
|
func (set *heldSessionSet) construct(ctx context.Context, topology heldSessionSetTopology) error {
|
|
acquisitionContext, cancelAcquisition := context.WithCancel(ctx)
|
|
defer cancelAcquisition()
|
|
ready := make(chan struct{}, len(set.plans))
|
|
release := make(chan struct{})
|
|
var releaseOnce sync.Once
|
|
releaseWorkers := func() { releaseOnce.Do(func() { close(release) }) }
|
|
group, groupContext := errgroup.WithContext(acquisitionContext)
|
|
var mu sync.Mutex
|
|
var constructionErrors error
|
|
for index, plan := range set.plans {
|
|
index, plan := index, plan
|
|
group.Go(func() error {
|
|
select {
|
|
case ready <- struct{}{}:
|
|
case <-groupContext.Done():
|
|
return groupContext.Err()
|
|
}
|
|
select {
|
|
case <-release:
|
|
case <-groupContext.Done():
|
|
return groupContext.Err()
|
|
}
|
|
topologyAgent := topology.agents[plan.Agent.Int()]
|
|
// Acquisition cancellation unblocks siblings; successful sessions retain the outer lifetime context.
|
|
session, err := constructHeldSession(groupContext, ctx, topology.dashboard, topologyAgent, plan, set.dependencies)
|
|
if err != nil {
|
|
mu.Lock()
|
|
constructionErrors = errors.Join(constructionErrors, err)
|
|
mu.Unlock()
|
|
return err
|
|
}
|
|
mu.Lock()
|
|
set.sessions[index] = session
|
|
mu.Unlock()
|
|
return nil
|
|
})
|
|
}
|
|
for range set.plans {
|
|
select {
|
|
case <-ready:
|
|
case <-groupContext.Done():
|
|
releaseWorkers()
|
|
groupError := group.Wait()
|
|
return errors.Join(groupContext.Err(), groupError, constructionErrors, set.rollback(ctx))
|
|
}
|
|
}
|
|
releaseWorkers()
|
|
groupError := group.Wait()
|
|
if groupError != nil {
|
|
return errors.Join(groupError, constructionErrors, set.rollback(ctx))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func constructHeldSession(ctx, lifetimeContext context.Context, dashboardInstance *dashboard.Dashboard, topology HeldSessionAgent, plan StressSessionPlan, dependencies HeldSessionSetDependencies) (heldSession, error) {
|
|
switch plan.Kind {
|
|
case StressSessionTerminal:
|
|
return dependencies.Terminal(ctx, heldTerminalInput{Dashboard: dashboardInstance, PATClient: topology.PATClient, Agent: topology.Agent, Readiness: topology.Readiness, Plan: plan, LifetimeContext: lifetimeContext})
|
|
case StressSessionNAT:
|
|
return dependencies.NAT(ctx, heldNATInput{Dashboard: dashboardInstance, PATClient: topology.PATClient, Agent: topology.Agent, Readiness: topology.Readiness, Plan: plan, LifetimeContext: lifetimeContext})
|
|
case StressSessionFM:
|
|
return dependencies.FM(ctx, heldLegacyFMInput{Dashboard: dashboardInstance, PATClient: topology.PATClient, Agent: topology.Agent, Readiness: topology.Readiness, Plan: plan, LifetimeContext: lifetimeContext})
|
|
default:
|
|
return nil, fmt.Errorf("unsupported held session kind %q: %w", plan.Kind, ErrInvalidHeldSessionSetPlan)
|
|
}
|
|
}
|
|
|
|
func (set *heldSessionSet) waitLive(ctx context.Context) error {
|
|
group, groupContext := errgroup.WithContext(ctx)
|
|
for index, session := range set.sessions {
|
|
index, session := index, session
|
|
group.Go(func() error {
|
|
if session == nil {
|
|
return fmt.Errorf("session %d was not constructed", index)
|
|
}
|
|
if err := session.WaitLive(groupContext); err != nil {
|
|
return err
|
|
}
|
|
select {
|
|
case <-session.Done():
|
|
return errors.Join(ErrHeldSessionPrematureClose, session.CloseResult())
|
|
default:
|
|
}
|
|
if !sessionMatchesPlan(session, set.plans[index]) {
|
|
return fmt.Errorf("session %d does not match its plan", index)
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
if err := group.Wait(); err != nil {
|
|
return err
|
|
}
|
|
for _, session := range set.sessions {
|
|
if session == nil {
|
|
return errors.New("held session set has an unconstructed session")
|
|
}
|
|
}
|
|
streamIDs, err := set.streamIDs()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
streamErr := waitHeldSessionSetStreams(ctx, set.state, set.baseline.Count+len(streamIDs), streamIDs, true, set.dependencies.WaitState)
|
|
_, aggregateErr := set.dependencies.WaitState(ctx, set.state, client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(set.baseline.Count + len(streamIDs))})
|
|
return errors.Join(streamErr, aggregateErr)
|
|
}
|
|
|
|
func (set *heldSessionSet) streamIDs() ([]string, error) {
|
|
ids := make([]string, len(set.sessions))
|
|
seen := make(map[string]struct{}, len(ids))
|
|
for index, session := range set.sessions {
|
|
streamID, present := session.IOStreamID()
|
|
if !present || streamID == "" {
|
|
return nil, errors.New("held session stream ID is empty")
|
|
}
|
|
if _, exists := seen[streamID]; exists {
|
|
return nil, fmt.Errorf("duplicate held session stream ID: %w", ErrInvalidHeldSessionSetPlan)
|
|
}
|
|
seen[streamID] = struct{}{}
|
|
ids[index] = streamID
|
|
}
|
|
return ids, nil
|
|
}
|
|
|
|
func sessionMatchesPlan(session heldSession, plan StressSessionPlan) bool {
|
|
return session.Plan() == plan
|
|
}
|
|
|
|
func waitHeldSessionSetStreams(ctx context.Context, state heldSessionSetStateObserver, expectedCount int, streamIDs []string, present bool, waitState func(context.Context, heldSessionSetStateObserver, client.IOStreamStateExpectation) (client.IOStreamState, error)) error {
|
|
group, groupContext := errgroup.WithContext(ctx)
|
|
var mu sync.Mutex
|
|
var joined error
|
|
for _, streamID := range streamIDs {
|
|
streamID := streamID
|
|
group.Go(func() error {
|
|
expectation := client.IOStreamStateExpectation{ExpectedCount: client.ExpectedIOStreamCount(expectedCount)}
|
|
if present {
|
|
expectation.PresentStreamID = streamID
|
|
} else {
|
|
expectation.AbsentStreamID = streamID
|
|
}
|
|
_, err := waitState(groupContext, state, expectation)
|
|
if err != nil {
|
|
mu.Lock()
|
|
joined = errors.Join(joined, err)
|
|
mu.Unlock()
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
return errors.Join(group.Wait(), joined)
|
|
}
|