Fix tsh player issues (#11491)

This commit fixes race conditions in the tsh session player by using
a condition variable to detect state changes rather than unsafely
polling a variable that is written by a separate goroutine.

In addition, fix an off by one error when resuming playback
after pausing. The player's position variable has always stored
the index of the last succesfully played event, so when we resume
playback we should start at position+1 to not re-play the previous
event twice.

Fixes #11479

lib/client/api.go

@@ -3153,7 +3153,7 @@ func playSession(sessionEvents []events.EventFields, stream []byte) error {
 	)
 	// playback control goroutine
 	go func() {
-		defer player.Stop()
+		defer player.RequestStop()
 		var key [1]byte
 		for {
 			_, err := term.Stdin().Read(key[:])

lib/client/player.go

@@ -25,10 +25,13 @@ import (
 
 	"github.com/gravitational/teleport/lib/client/terminal"
 	"github.com/gravitational/teleport/lib/events"
+	"github.com/jonboulle/clockwork"
 )
 
+type tshPlayerState int
+
 const (
-	stateStopped = iota
+	stateStopped tshPlayerState = iota
 	stateStopping
 	statePlaying
 )
@@ -37,60 +40,67 @@ const (
 // and allows to control it
 type sessionPlayer struct {
 	sync.Mutex
+	cond *sync.Cond
+
+	state    tshPlayerState
+	position int // position is the index of the last event successfully played back
+
+	clock         clockwork.Clock
 	stream        []byte
 	sessionEvents []events.EventFields
 	term          *terminal.Terminal
 
-	state    int
-	position int
-
 	// stopC is used to tell the caller that player has finished playing
-	stopC chan int
+	stopC    chan int
+	stopOnce sync.Once
 }
 
 func newSessionPlayer(sessionEvents []events.EventFields, stream []byte, term *terminal.Terminal) *sessionPlayer {
-	return &sessionPlayer{
+	p := &sessionPlayer{
+		clock:         clockwork.NewRealClock(),
+		position:      -1, // position is the last successfully written event
 		stream:        stream,
 		sessionEvents: sessionEvents,
-		stopC:         make(chan int),
 		term:          term,
+		stopC:         make(chan int),
 	}
+	p.cond = sync.NewCond(p)
+	return p
 }
 
 func (p *sessionPlayer) Play() {
 	p.playRange(0, 0)
 }
 
-func (p *sessionPlayer) Stop() {
+func (p *sessionPlayer) Stopped() bool {
 	p.Lock()
 	defer p.Unlock()
-	if p.stopC != nil {
-		close(p.stopC)
-		p.stopC = nil
-	}
-}
-
-func (p *sessionPlayer) Stopped() bool {
 	return p.state == stateStopped
 }
 
 func (p *sessionPlayer) Rewind() {
 	p.Lock()
 	defer p.Unlock()
 	if p.state != stateStopped {
-		p.state = stateStopping
+		p.setState(stateStopping)
 		p.waitUntil(stateStopped)
 	}
 	if p.position > 0 {
 		p.playRange(p.position-1, p.position)
 	}
 }
 
+func (p *sessionPlayer) stopRequested() bool {
+	p.Lock()
+	defer p.Unlock()
+	return p.state == stateStopping
+}
+
 func (p *sessionPlayer) Forward() {
 	p.Lock()
 	defer p.Unlock()
 	if p.state != stateStopped {
-		p.state = stateStopping
+		p.setState(stateStopping)
 		p.waitUntil(stateStopped)
 	}
 	if p.position < len(p.sessionEvents) {
@@ -102,20 +112,44 @@ func (p *sessionPlayer) TogglePause() {
 	p.Lock()
 	defer p.Unlock()
 	if p.state == statePlaying {
-		p.state = stateStopping
+		p.setState(stateStopping)
 		p.waitUntil(stateStopped)
 	} else {
-		p.playRange(p.position, 0)
+		p.playRange(p.position+1, 0)
 		p.waitUntil(statePlaying)
 	}
 }
 
-func (p *sessionPlayer) waitUntil(state int) {
+// RequestStop makes an asynchronous request for the player to stop playing.
+// Playback may not stop before this method returns.
+func (p *sessionPlayer) RequestStop() {
+	p.Lock()
+	defer p.Unlock()
+
+	switch p.state {
+	case stateStopped, stateStopping:
+		// do nothing if stop already in progress
+	default:
+		p.setState(stateStopping)
+	}
+}
+
+// waitUntil waits for the specified state to be reached.
+// Callers must hold the lock on p.Mutex before calling.
+func (p *sessionPlayer) waitUntil(state tshPlayerState) {
 	for state != p.state {
-		time.Sleep(time.Millisecond)
+		p.cond.Wait()
 	}
 }
 
+// setState sets the current player state and notifies any
+// goroutines waiting in waitUntil(). Callers must hold the
+// lock on p.Mutex before calling.
+func (p *sessionPlayer) setState(state tshPlayerState) {
+	p.state = state
+	p.cond.Broadcast()
+}
+
 // timestampFrame prints 'event timestamp' in the top right corner of the
 // terminal after playing every 'print' event
 func timestampFrame(term *terminal.Terminal, message string) {
@@ -146,56 +180,52 @@ func timestampFrame(term *terminal.Terminal, message string) {
 // that playback starts from there.
 func (p *sessionPlayer) playRange(from, to int) {
 	if to > len(p.sessionEvents) || from < 0 {
-		p.state = stateStopped
+		p.Lock()
+		p.setState(stateStopped)
+		p.Unlock()
 		return
 	}
 	if to == 0 {
 		to = len(p.sessionEvents)
 	}
 	// clear screen between runs:
 	os.Stdout.Write([]byte("\x1bc"))
-	// wait: waits between events during playback
-	prev := time.Duration(0)
-	wait := func(i int, e events.EventFields) {
-		ms := time.Duration(e.GetInt("ms"))
-		// before "from"? play that instantly:
-		if i >= from {
-			delay := ms - prev
-			// make playback smoother:
-			if delay < 10 {
-				delay = 0
-			}
-			if delay > 250 && delay < 500 {
-				delay = 250
-			}
-			if delay > 500 && delay < 1000 {
-				delay = 500
-			}
-			if delay > 1000 {
-				delay = 1000
-			}
-			timestampFrame(p.term, e.GetString("time"))
-			time.Sleep(time.Millisecond * delay)
-		}
-		prev = ms
-	}
+
 	// playback goroutine:
 	go func() {
+		var i int
+
 		defer func() {
-			p.state = stateStopped
+			p.Lock()
+			p.setState(stateStopped)
+			p.Unlock()
+
+			// played last event?
+			if i == len(p.sessionEvents) {
+				p.stopOnce.Do(func() { close(p.stopC) })
+			}
 		}()
-		p.state = statePlaying
-		i, offset, bytes := 0, 0, 0
+
+		p.Lock()
+		p.setState(statePlaying)
+		p.Unlock()
+
+		prev := time.Duration(0)
+		offset, bytes := 0, 0
 		for i = 0; i < to; i++ {
-			if p.state == stateStopping {
+			if p.stopRequested() {
 				return
 			}
+
 			e := p.sessionEvents[i]
 
 			switch e.GetString(events.EventType) {
 			// 'print' event (output)
 			case events.SessionPrintEvent:
-				wait(i, e)
+				// delay is only necessary once we've caught up to the "from" event
+				if i >= from {
+					prev = p.applyDelay(prev, e)
+				}
 				offset = e.GetInt("offset")
 				bytes = e.GetInt("bytes")
 				os.Stdout.Write(p.stream[offset : offset+bytes])
@@ -211,11 +241,32 @@ func (p *sessionPlayer) playRange(from, to int) {
 			default:
 				continue
 			}
+			p.Lock()
 			p.position = i
-		}
-		// played last event?
-		if i == len(p.sessionEvents) {
-			p.Stop()
+			p.Unlock()
 		}
 	}()
 }
+
+// applyDelay waits until it is time to play back the current event.
+// It returns the duration from the start of the session up until the current event.
+func (p *sessionPlayer) applyDelay(previousTimestamp time.Duration, e events.EventFields) time.Duration {
+	eventTime := time.Duration(e.GetInt("ms") * int(time.Millisecond))
+	delay := eventTime - previousTimestamp
+
+	// make playback smoother:
+	switch {
+	case delay < 10*time.Millisecond:
+		delay = 0
+	case delay > 250*time.Millisecond && delay < 500*time.Millisecond:
+		delay = 250 * time.Millisecond
+	case delay > 500*time.Millisecond && delay < 1*time.Second:
+		delay = 500 * time.Millisecond
+	case delay > time.Second:
+		delay = time.Second
+	}
+
+	timestampFrame(p.term, e.GetString("time"))
+	p.clock.Sleep(delay)
+	return eventTime
+}