Merge pull request #225 from walles/johan/select-shutdown

Fancy shutdown

twin/screen-setup-windows.go

@@ -5,13 +5,60 @@ package twin
 
 import (
 	"fmt"
+	"io"
 	"os"
+	"sync/atomic"
 	"syscall"
 
 	"golang.org/x/sys/windows"
 	"golang.org/x/term"
 )
 
+type interruptableReaderImpl struct {
+	base              *os.File
+	shutdownRequested atomic.Bool
+}
+
+// NOTE: To work properly, this Read() should return immediately after somebody
+// calls Interrupt(), *without first reading any bytes from the base reader*.
+//
+// This implementation doesn't do that. If you want to fix this, the not-Windows
+// implementation in screen-setup.go may or may not work as inspiration.
+func (r *interruptableReaderImpl) Read(p []byte) (n int, err error) {
+	if r.shutdownRequested.Load() {
+		err = io.EOF
+		return
+	}
+
+	n, err = r.base.Read(p)
+	if err != nil {
+		return
+	}
+
+	if r.shutdownRequested.Load() {
+		err = io.EOF
+	}
+	return
+}
+
+func (r *interruptableReaderImpl) Interrupt() {
+	// Previously we used to close the screen.ttyIn file descriptor here, but:
+	// * That didn't interrupt the blocking read() in the main loop
+	// * It may or may not have caused shutdown issues on Windows
+	//
+	// Setting this flag doesn't interrupt the blocking read() either, but it
+	// should at least not cause any shutdown issues on Windows.
+	//
+	// Ref:
+	// * https://github.com/walles/moar/issues/217
+	// * https://github.com/walles/moar/issues/221
+	r.shutdownRequested.Store(true)
+}
+
+func newInterruptableReader(base *os.File) (interruptableReader, error) {
+	return &interruptableReaderImpl{base: base}, nil
+}
+
 func (screen *UnixScreen) setupSigwinchNotification() {
 	screen.sigwinch = make(chan int, 1)
 	screen.sigwinch <- 0 // Trigger initial screen size query

twin/screen-setup.go

@@ -4,14 +4,95 @@
 package twin
 
 import (
+	"io"
 	"os"
 	"os/signal"
 	"syscall"
 
 	log "github.com/sirupsen/logrus"
+	"golang.org/x/sys/unix"
 	"golang.org/x/term"
 )
 
+type interruptableReaderImpl struct {
+	base *os.File
+
+	shutdownPipeReader *os.File
+	shutdownPipeWriter *os.File
+
+	interruptionComplete chan struct{}
+}
+
+func (r *interruptableReaderImpl) Read(p []byte) (n int, err error) {
+	// "This argument should be set to the highest-numbered file descriptor in
+	// any of the three sets, plus 1. The indicated file descriptors in each set
+	// are checked, up to this limit"
+	//
+	// Ref: https://man7.org/linux/man-pages/man2/select.2.html
+	nfds := r.base.Fd()
+	if r.shutdownPipeReader.Fd() > nfds {
+		nfds = r.shutdownPipeReader.Fd()
+	}
+
+	readFds := unix.FdSet{}
+	readFds.Set(int(r.shutdownPipeReader.Fd()))
+	readFds.Set(int(r.base.Fd()))
+
+	_, err = unix.Select(int(nfds)+1, &readFds, nil, nil, nil)
+	if err != nil {
+		// Select failed
+		return
+	}
+
+	if readFds.IsSet(int(r.shutdownPipeReader.Fd())) {
+		// Shutdown requested
+		closeErr := r.shutdownPipeReader.Close()
+		if closeErr != nil {
+			// This should never happen, but if it does we should log it
+			log.Debug("Failed to close shutdown pipe reader: ", closeErr)
+		}
+
+		err = io.EOF
+
+		// Let Interrupt() know we're done
+		r.interruptionComplete <- struct{}{}
+
+		return
+	}
+
+	if readFds.IsSet(int(r.base.Fd())) {
+		// Base has stuff
+		return r.base.Read(p)
+	}
+
+	// Neither base nor shutdown pipe was ready, this should never happen
+	return
+}
+
+func (r *interruptableReaderImpl) Interrupt() {
+	// This will make the select() call claim the read end is ready
+	r.shutdownPipeWriter.Close()
+
+	// Wait for reader to react
+	<-r.interruptionComplete
+}
+
+func newInterruptableReader(base *os.File) (interruptableReader, error) {
+	reader := interruptableReaderImpl{
+		base:                 base,
+		interruptionComplete: make(chan struct{}),
+	}
+	pr, pw, err := os.Pipe()
+	if err != nil {
+		return nil, err
+	}
+
+	reader.shutdownPipeReader = pr
+	reader.shutdownPipeWriter = pw
+
+	return &reader, nil
+}
+
 func (screen *UnixScreen) setupSigwinchNotification() {
 	screen.sigwinch = make(chan int, 1)
 	screen.sigwinch <- 0 // Trigger initial screen size query

twin/screen.go

@@ -7,7 +7,6 @@ import (
 	"regexp"
 	"strconv"
 	"strings"
-	"sync/atomic"
 	"unicode/utf8"
 
 	log "github.com/sirupsen/logrus"
@@ -70,6 +69,13 @@ type Screen interface {
 	Events() chan Event
 }
 
+type interruptableReader interface {
+	Read(p []byte) (n int, err error)
+
+	// Interrupt unblocks the read call, either now or eventually.
+	Interrupt()
+}
+
 type UnixScreen struct {
 	widthAccessFromSizeOnly  int // Access from Size() method only
 	heightAccessFromSizeOnly int // Access from Size() method only
@@ -81,7 +87,7 @@ type UnixScreen struct {
 
 	events chan Event
 
-	shutdownRequested atomic.Bool
+	ttyInReader interruptableReader
 
 	ttyIn            *os.File
 	oldTerminalState *term.State //nolint Not used on Windows
@@ -150,6 +156,14 @@ func NewScreenWithMouseModeAndColorType(mouseMode MouseMode, terminalColorCount
 	if err != nil {
 		return nil, fmt.Errorf("problem setting up TTY: %w", err)
 	}
+	screen.ttyInReader, err = newInterruptableReader(screen.ttyIn)
+	if err != nil {
+		restoreErr := screen.restoreTtyInTtyOut()
+		if restoreErr != nil {
+			log.Warn("Problem restoring TTY state after failed interruptable reader setup: ", restoreErr)
+		}
+		return nil, fmt.Errorf("problem setting up TTY reader: %w", err)
+	}
 
 	screen.setAlternateScreenMode(true)
 
@@ -176,18 +190,8 @@ func (screen *UnixScreen) Close() {
 	// Tell the pager to exit unless it hasn't already
 	screen.events <- EventExit{}
 
-	// Tell our main loop to exit. Previously we used to close the screen.ttyIn
-	// file descriptor here, but:
-	// * That didn't interrupt the blocking read() in the main loop
-	// * It may or may not have caused shutdown issues on Windows
-	//
-	// Setting this flag doesn't interrupt the blocking read() either, but it
-	// is should at least be less likely to cause shutdown issues on Windows.
-	//
-	// Ref:
-	// * https://github.com/walles/moar/issues/217
-	// * https://github.com/walles/moar/issues/221
-	screen.shutdownRequested.Store(true)
+	// Tell our main loop to exit
+	screen.ttyInReader.Interrupt()
 
 	screen.hideCursor(false)
 	screen.enableMouseTracking(false)
@@ -375,7 +379,7 @@ func (screen *UnixScreen) mainLoop() {
 	maxBytesRead := 0
 	expectingTerminalBackgroundColor := true
 	for {
-		count, err := screen.ttyIn.Read(buffer)
+		count, err := screen.ttyInReader.Read(buffer)
 		if err != nil {
 			// Ref:
 			// * https://github.com/walles/moar/issues/145
@@ -386,12 +390,6 @@ func (screen *UnixScreen) mainLoop() {
 			screen.events <- EventExit{}
 			return
 		}
-		if screen.shutdownRequested.Load() {
-			log.Debug("Shutdown requested, twin giving up")
-
-			screen.events <- EventExit{}
-			return
-		}
 
 		if expectingTerminalBackgroundColor {
 			// This is the response to our background color request