1//
2// Copyright (c) 2018, Joyent, Inc. All rights reserved.
3//
4// This Source Code Form is subject to the terms of the Mozilla Public
5// License, v. 2.0. If a copy of the MPL was not distributed with this
6// file, You can obtain one at http://mozilla.org/MPL/2.0/.
7//
8
9package testutils
10
11import (
12	"log"
13	"sync"
14	"sync/atomic"
15)
16
17type runState int32
18
19const (
20	stateIdle runState = iota
21	stateRunning
22	stateCancelling
23)
24
25type basicRunner struct {
26	Steps []Step
27
28	cancelCh chan struct{}
29	doneCh   chan struct{}
30	state    runState
31	l        sync.Mutex
32}
33
34func (b *basicRunner) Run(state TritonStateBag) {
35	b.l.Lock()
36	if b.state != stateIdle {
37		panic("already running")
38	}
39
40	cancelCh := make(chan struct{})
41	doneCh := make(chan struct{})
42	b.cancelCh = cancelCh
43	b.doneCh = doneCh
44	b.state = stateRunning
45	b.l.Unlock()
46
47	defer func() {
48		b.l.Lock()
49		b.cancelCh = nil
50		b.doneCh = nil
51		b.state = stateIdle
52		close(doneCh)
53		b.l.Unlock()
54	}()
55
56	// This goroutine listens for cancels and puts the StateCancelled key
57	// as quickly as possible into the state bag to mark it.
58	go func() {
59		select {
60		case <-cancelCh:
61			// Flag cancel and wait for finish
62			state.Put(StateCancelled, true)
63			<-doneCh
64		case <-doneCh:
65		}
66	}()
67
68	for _, step := range b.Steps {
69		// We also check for cancellation here since we can't be sure
70		// the goroutine that is running to set it actually ran.
71		if runState(atomic.LoadInt32((*int32)(&b.state))) == stateCancelling {
72			state.Put(StateCancelled, true)
73			break
74		}
75
76		action := step.Run(state)
77		defer step.Cleanup(state)
78
79		if _, ok := state.GetOk(StateCancelled); ok {
80			break
81		}
82
83		if action == Halt {
84			log.Println("[INFO] Halt requested by current step")
85			state.Put(StateHalted, true)
86			break
87		}
88	}
89}
90
91func (b *basicRunner) Cancel() {
92	b.l.Lock()
93	switch b.state {
94	case stateIdle:
95		// Not running, so Cancel is... done.
96		b.l.Unlock()
97		return
98	case stateRunning:
99		// Running, so mark that we cancelled and set the state
100		close(b.cancelCh)
101		b.state = stateCancelling
102		fallthrough
103	case stateCancelling:
104		// Already cancelling, so just wait until we're done
105		ch := b.doneCh
106		b.l.Unlock()
107		<-ch
108	}
109}
110