Skip to content

Commit 3219298

Browse files
authored
fix(control): refactor the controller to remove data races (#939)
This refactors the controller query and how the controller interacts with that query. Previously, the query had an ill-defined lifecycle and it could be arbitrarily canceled at any time that caused race conditions when transitioning between states. The transition touched values that were not protected by a mutex which caused a data race. The query has now been changed so the `Cancel()` method only _signals_ that the query should be canceled and does that by using the query context and sending that downstream. This causes the state to transition lazily so when the query attempts to transition to a different state, it will automatically transition to the canceled state and report that it failed to transition to a new state. When the query controller attempts to report an error, it realizes that the query has been canceled and uses the error on the query context instead. In addition, to ensure that the controller reports the correct state of a query to a task manager, the `State()` method has been modified to check the context to see if the query has been canceled even if the stateh hasn't been updated. With these changes, it means that all state transitions and changes to sensitive state happen within the controller execution loop instead of potentially running outside of it. The channel returned by `Ready()` will only be set when an error is caught during the execution lifecycle (such as by cancellation) or when the query has successfully executed. This makes the state machine much more simple than previously.
1 parent 32b160b commit 3219298

2 files changed

Lines changed: 301 additions & 79 deletions

File tree

‎control/controller.go‎

Lines changed: 147 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -112,14 +112,14 @@ func New(c Config) *Controller {
112112
func (c *Controller) Query(ctx context.Context, compiler flux.Compiler) (flux.Query, error) {
113113
q := c.createQuery(ctx, compiler.CompilerType())
114114
if err := c.compileQuery(q, compiler); err != nil {
115-
q.transitionTo(Errored)
115+
q.setErr(err)
116116
c.countQueryRequest(q, labelCompileError)
117-
return nil, err
117+
return nil, q.Err()
118118
}
119119
if err := c.enqueueQuery(q); err != nil {
120-
q.transitionTo(Errored)
120+
q.setErr(err)
121121
c.countQueryRequest(q, labelQueueError)
122-
return nil, err
122+
return nil, q.Err()
123123
}
124124
c.countQueryRequest(q, labelSuccess)
125125
return q, nil
@@ -154,15 +154,14 @@ func (c *Controller) createQuery(ctx context.Context, ct flux.CompilerType) *Que
154154
c.metrics.allDur.WithLabelValues(labelValues...),
155155
c.metrics.all.WithLabelValues(labelValues...),
156156
)
157-
ready := make(chan map[string]flux.Result, 1)
158157
return &Query{
159158
id: id,
160159
labelValues: labelValues,
161160
compileLabelValues: compileLabelValues,
162161
state: Created,
163162
c: c,
164163
now: time.Now().UTC(),
165-
ready: ready,
164+
ready: make(chan map[string]flux.Result, 1),
166165
parentCtx: parentCtx,
167166
parentSpan: parentSpan,
168167
cancel: cancel,
@@ -346,7 +345,7 @@ func (c *Controller) run() {
346345
pq.Pop()
347346
}
348347
if err != nil {
349-
go q.setErr(err)
348+
q.setErr(err)
350349
}
351350
}
352351
}
@@ -471,13 +470,13 @@ type Query struct {
471470
spec flux.Spec
472471
now time.Time
473472

474-
err error
475-
476473
ready chan map[string]flux.Result
477474

478-
mu sync.Mutex
479-
state State
480-
cancel func()
475+
// query state. The stateMu protects access for the group below.
476+
stateMu sync.RWMutex
477+
state State
478+
err error
479+
cancel func()
481480

482481
parentCtx, currentCtx context.Context
483482
parentSpan, currentSpan *span
@@ -508,52 +507,48 @@ func (q *Query) Concurrency() int {
508507

509508
// Cancel will stop the query execution.
510509
func (q *Query) Cancel() {
511-
q.mu.Lock()
512-
defer q.mu.Unlock()
513-
514-
// call cancel func
510+
// Call the cancel function to signal that execution should
511+
// be interrupted.
515512
q.cancel()
516-
517-
if q.state != Errored && q.state != Canceled {
518-
q.transitionTo(Canceled)
519-
}
520513
}
521514

522515
// Ready returns a channel that will deliver the query results.
516+
//
523517
// It's possible that the channel is closed before any results arrive.
524-
// In particular, if a query's context is cancelled during execution,
525-
// the channel will be closed immediately and an error will be exposed
526-
// via Err(). For this reason, a query should always be inspected for
527-
// an error using Err().
518+
// In particular, if a query's context or the query itself is canceled,
519+
// the query may close the results channel before any results are computed.
520+
//
521+
// The query may also have an error during execution so the Err()
522+
// function should be used to check if an error happened.
528523
func (q *Query) Ready() <-chan map[string]flux.Result {
529524
return q.ready
530525
}
531526

532-
// Done must always be called to free resources.
527+
// Done signals to the Controller that this query is no longer
528+
// being used and resources related to the query may be freed.
529+
//
530+
// The Ready method must have returned a result before calling
531+
// this method either by the query executing, being canceled, or
532+
// an error occurring.
533533
func (q *Query) Done() {
534-
// TODO(jsternberg): I'm pretty sure this can results in releasing resources that are still
535-
// actively being used. This method should indicate to the controller that it no longer has to
536-
// keep the queries around, but it shouldn't let the client decide the resources are freed since
537-
// the query may have been canceled and is still waiting to hit a cancel point.
538-
// This could result in reporting the wrong number for resource usage.
534+
// We are not considered to be in the run loop anymore once
535+
// this is called.
539536
q.done.Do(func() {
540-
q.mu.Lock()
541-
if q.state != Errored {
542-
q.transitionTo(Finished)
543-
}
544-
q.mu.Unlock()
537+
q.stateMu.Lock()
538+
q.transitionTo(Finished)
539+
q.stateMu.Unlock()
545540

546-
// We must send the information that this query is done outside of the lock because
547-
// otherwise we can deadlock with the controller.
548541
q.c.queryDone <- q
549542
})
550543
}
551544

552-
// Statistics reports the statisitcs for the query.
553-
// The statisitcs are not complete until the query is finished.
545+
// Statistics reports the statistics for the query.
546+
//
547+
// This method must be called after Done. It will block until
548+
// the query has been finalized unless a context is given.
554549
func (q *Query) Statistics() flux.Statistics {
555-
q.mu.Lock()
556-
defer q.mu.Unlock()
550+
q.stateMu.Lock()
551+
defer q.stateMu.Unlock()
557552

558553
stats := q.stats
559554
stats.Concurrency = q.concurrency
@@ -565,17 +560,45 @@ func (q *Query) Statistics() flux.Statistics {
565560

566561
// State reports the current state of the query.
567562
func (q *Query) State() State {
568-
q.mu.Lock()
569-
s := q.state
570-
q.mu.Unlock()
571-
return s
563+
q.stateMu.RLock()
564+
state := q.state
565+
if !isFinishedState(state) {
566+
// If the query is a non-finished state, check the
567+
// context to see if we have been interrupted.
568+
select {
569+
case <-q.parentCtx.Done():
570+
// The query has been canceled so report to the
571+
// outside world that we have been canceled.
572+
// Do NOT attempt to change the internal state
573+
// variable here. It is a minefield. Leave the
574+
// normal query execution to figure that out.
575+
state = Canceled
576+
default:
577+
// The context has not been canceled.
578+
}
579+
}
580+
q.stateMu.RUnlock()
581+
return state
572582
}
573583

574584
// transitionTo will transition from one state to another. If a list of current states
575585
// is given, then the query must be in one of those states for the transition to succeed.
576-
// This method must be called with a lock and it must not transition to a final state
577-
// from within the run loop.
586+
// This method must be called with a lock and it must be called from within the run loop.
578587
func (q *Query) transitionTo(newState State, currentState ...State) bool {
588+
// If we are transitioning to a non-finished state, the query
589+
// may have been canceled. If the query was canceled, then
590+
// we need to transition to the canceled state
591+
if !isFinishedState(newState) {
592+
select {
593+
case <-q.parentCtx.Done():
594+
// Transition to the canceled state and report that
595+
// we failed to transition to the desired state.
596+
_ = q.transitionTo(Canceled)
597+
return false
598+
default:
599+
}
600+
}
601+
579602
if len(currentState) > 0 {
580603
// Find the current state in the list of current states.
581604
for _, st := range currentState {
@@ -606,17 +629,11 @@ TRANSITION:
606629
q.currentSpan, q.currentCtx = nil, nil
607630

608631
// If we are transitioning to a finished state from a non-finished state, finish the parent span.
609-
switch newState {
610-
case Errored, Canceled, Finished:
632+
if isFinishedState(newState) {
611633
if q.parentSpan != nil {
612634
q.parentSpan.Finish()
613635
q.stats.TotalDuration = q.parentSpan.Duration
614636
q.parentSpan = nil
615-
616-
// Close the ready channel on the first time we move to one of these states.
617-
// It should signal any queries waiting on the results that no results will come.
618-
// Signal to the main loop that this query has completed.
619-
close(q.ready)
620637
}
621638
}
622639

@@ -655,63 +672,81 @@ TRANSITION:
655672
}
656673

657674
func (q *Query) isOK() bool {
658-
q.mu.Lock()
675+
q.stateMu.RLock()
659676
ok := q.state != Canceled && q.state != Errored && q.state != Finished
660-
q.mu.Unlock()
677+
q.stateMu.RUnlock()
661678
return ok
662679
}
663680

664681
// Err reports any error the query may have encountered.
665682
func (q *Query) Err() error {
666-
q.mu.Lock()
683+
q.stateMu.Lock()
667684
err := q.err
668-
q.mu.Unlock()
685+
q.stateMu.Unlock()
669686
return err
670687
}
688+
689+
// setErr marks this query with an error. If the query was
690+
// canceled, then the error is ignored.
691+
//
692+
// This will mark the query as ready so setResults must not
693+
// be called if this method is invoked.
671694
func (q *Query) setErr(err error) {
672-
q.mu.Lock()
673-
defer q.mu.Unlock()
695+
q.stateMu.Lock()
696+
defer q.stateMu.Unlock()
697+
698+
// We may have this get called when the query is canceled.
699+
// If that is the case, transition to the canceled state
700+
// instead and record the error from that since the error
701+
// we received is probably wrong.
702+
select {
703+
case <-q.parentCtx.Done():
704+
q.transitionTo(Canceled)
705+
err = q.parentCtx.Err()
706+
default:
707+
q.transitionTo(Errored)
708+
}
674709
q.err = err
675710

676-
q.transitionTo(Errored)
711+
// Close the ready channel to report that no results
712+
// will be sent.
713+
close(q.ready)
677714
}
678715

716+
// setResults will set the results and send them over the ready channel.
679717
func (q *Query) setResults(r map[string]flux.Result) {
680-
q.mu.Lock()
681-
if q.state == Executing {
682-
q.ready <- r
683-
}
684-
q.mu.Unlock()
718+
q.ready <- r
719+
close(q.ready)
685720
}
686721

687722
// tryCompile attempts to transition the query into the Compiling state.
688723
func (q *Query) tryCompile() bool {
689-
q.mu.Lock()
690-
defer q.mu.Unlock()
724+
q.stateMu.Lock()
725+
defer q.stateMu.Unlock()
691726

692727
return q.transitionTo(Compiling, Created)
693728
}
694729

695730
// tryPlan attempts to transition the query into the Planning state.
696731
func (q *Query) tryPlan() bool {
697-
q.mu.Lock()
698-
defer q.mu.Unlock()
732+
q.stateMu.Lock()
733+
defer q.stateMu.Unlock()
699734

700735
return q.transitionTo(Planning, Compiling)
701736
}
702737

703738
// tryQueue attempts to transition the query into the Queueing state.
704739
func (q *Query) tryQueue() bool {
705-
q.mu.Lock()
706-
defer q.mu.Unlock()
740+
q.stateMu.Lock()
741+
defer q.stateMu.Unlock()
707742

708743
return q.transitionTo(Queueing, Planning)
709744
}
710745

711746
// tryRequeue attempts to transition the query into the Requeueing state.
712747
func (q *Query) tryRequeue() bool {
713-
q.mu.Lock()
714-
defer q.mu.Unlock()
748+
q.stateMu.Lock()
749+
defer q.stateMu.Unlock()
715750

716751
if q.state == Requeueing {
717752
// Already in the correct state.
@@ -722,8 +757,8 @@ func (q *Query) tryRequeue() bool {
722757

723758
// tryExec attempts to transition the query into the Executing state.
724759
func (q *Query) tryExec() bool {
725-
q.mu.Lock()
726-
defer q.mu.Unlock()
760+
q.stateMu.Lock()
761+
defer q.stateMu.Unlock()
727762

728763
return q.transitionTo(Executing, Requeueing, Queueing)
729764
}
@@ -732,14 +767,40 @@ func (q *Query) tryExec() bool {
732767
type State int
733768

734769
const (
770+
// Created indicates the query has been created.
735771
Created State = iota
772+
773+
// Compiling indicates that the query is in the process
774+
// of executing the compiler associated with the query.
736775
Compiling
737-
Queueing
776+
777+
// Planning indicates that a query spec has been created
778+
// from the compiler and the query planner is executing.
738779
Planning
780+
781+
// Queueing indicates the query is waiting inside of the
782+
// scheduler to be executed.
783+
Queueing
784+
785+
// Requeueing indicates that the query scheduler wanted
786+
// to run the query, but not enough resources were available
787+
// so it is in the process of waiting again.
739788
Requeueing
789+
790+
// Executing indicates that the query is currently executing.
740791
Executing
792+
793+
// Errored indicates that there was an error when attempting
794+
// to execute a query within any state inside of the controller.
741795
Errored
796+
797+
// Finished indicates that the query has been marked as Done
798+
// and it is awaiting removal from the Controller or has already
799+
// been removed.
742800
Finished
801+
802+
// Canceled indicates that the query was signaled to be
803+
// canceled. A canceled query must still be released with Done.
743804
Canceled
744805
)
745806

@@ -768,6 +829,15 @@ func (s State) String() string {
768829
}
769830
}
770831

832+
func isFinishedState(state State) bool {
833+
switch state {
834+
case Canceled, Errored, Finished:
835+
return true
836+
default:
837+
return false
838+
}
839+
}
840+
771841
// span is a simple wrapper around opentracing.Span in order to
772842
// get access to the duration of the span for metrics reporting.
773843
type span struct {

0 commit comments

Comments
 (0)