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