Reapply "Read from standard input"

This reverts commit 88dee6a.

cmd/shui/main.go

@@ -80,6 +80,7 @@ var (
 	showStatus            bool     // Whether or not to show statuses.
 	showVersion           bool     // Whether or not to display a version.
 	silent                bool     // Whether or not to be quiet.
+	readStdin             bool     // Whether input comes from standard input.
 
 	// Supplied by the linker at build time.
 	version string
@@ -156,11 +157,21 @@ func main() {
 		os.Exit(0)
 	}
 
-	// Check for required flags.
-	if len(queries) == 0 {
-		flag.Usage()
-		fmt.Fprintf(os.Stderr, "Missing required argument -query\n")
-		os.Exit(1)
+	// Detect if running from standard input.
+	f, err := os.Stdin.Stat()
+	if err != nil {
+		panic(err)
+	}
+	if f.Mode()&os.ModeNamedPipe != 0 {
+		// We are reading standard input.
+		readStdin = true
+	} else {
+		// There is no standard input--queries are needed.
+		if len(queries) == 0 {
+			flag.Usage()
+			fmt.Fprintf(os.Stderr, "Missing required argument -query\n")
+			os.Exit(1)
+		}
 	}
 
 	// Set-up logging.
@@ -207,6 +218,7 @@ func main() {
 		PrometheusExporterAddr: promExporterAddr,
 		PushgatewayAddr:        promPushgatewayAddr,
 		Queries:                queries,
+		ReadStdin:              readStdin,
 	}
 
 	// Build display configuration.

internal/lib/config.go

@@ -19,7 +19,7 @@ type Config struct {
 	Count, Delay, DisplayMode, Mode                                                 int
 	ElasticsearchAddr, ElasticsearchIndex, ElasticsearchPassword, ElasticsearchUser string
 	Expressions, Filters, Labels, Queries                                           []string
-	History, LogMulti, Silent                                                       bool
+	History, LogMulti, ReadStdin, Silent                                            bool
 	LogLevel                                                                        string
 	Port                                                                            string
 	PrometheusExporterAddr                                                          string

internal/lib/query.go

@@ -4,8 +4,10 @@
 package lib
 
 import (
+	"bufio"
 	"io"
 	"log/slog"
+	"os"
 	"os/exec"
 	"strconv"
 	"time"
@@ -14,6 +16,11 @@ import (
 const (
 	QUERY_MODE_COMMAND int = iota + 1 // Queries are commands.
 	QUERY_MODE_PROFILE                // Queries are PIDs to profile.
+	QUERY_MODE_STDIN                  // Results are fron stdin.
+)
+
+var (
+	stdinScanner = bufio.NewScanner(os.Stdin) // Scanner for standard input queries.
 )
 
 // Wrapper for query execution.
@@ -90,6 +97,17 @@ func runQueryProfile(pid string, history bool) {
 	AddResult(pid, runProfile(pidInt), history)
 }
 
+// Reads standard input for results.
+func runQueryStdin(query string, history bool) {
+	slog.Debug("Reading stdin")
+
+	stdinScanner.Scan()
+	err := stdinScanner.Err()
+	e(err)
+
+	AddResult(query, stdinScanner.Text(), history)
+}
+
 // Entrypoint for 'query' mode.
 func Query(
 	queryMode, attempts, delay int,
@@ -107,21 +125,20 @@ func Query(
 	// Start the RPC server.
 	initServer(port)
 
-	for _, query := range queries {
-		// Initialize pause channels.
-		pauseQueryChans[query] = make(chan bool)
-	}
-
 	go func() {
 		// Wait for result consumption to become ready.
 		slog.Debug("Waiting for results readiness")
 		<-resultsReadyChan
 
-		for _, query := range queries {
-			// Execute the queries.
-			switch queryMode {
-			case QUERY_MODE_COMMAND:
-				slog.Debug("Executing in query mode command")
+		// Execute the queries.
+		switch queryMode {
+		case QUERY_MODE_COMMAND:
+			slog.Debug("Executing in query mode command")
+
+			for _, query := range queries {
+				// Initialize pause channels.
+				pauseQueryChans[query] = make(chan bool)
+
 				go runQuery(
 					query,
 					attempts,
@@ -131,8 +148,14 @@ func Query(
 					pauseQueryChans[query],
 					runQueryExec,
 				)
-			case QUERY_MODE_PROFILE:
-				slog.Debug("Executing in query mode profile")
+			}
+		case QUERY_MODE_PROFILE:
+			slog.Debug("Executing in query mode profile")
+
+			for _, query := range queries {
+				// Initialize pause channels.
+				pauseQueryChans[query] = make(chan bool)
+
 				go runQuery(
 					query,
 					attempts,
@@ -143,6 +166,22 @@ func Query(
 					runQueryProfile,
 				)
 			}
+		case QUERY_MODE_STDIN:
+			// When executing by reading standard input, there is only ever one "query".
+			slog.Debug("Executing in query mode stdin")
+
+			// Initialize pause channels.
+			pauseQueryChans[queries[0]] = make(chan bool)
+
+			go runQuery(
+				queries[0],
+				attempts,
+				delay,
+				history,
+				doneQueryChan,
+				pauseQueryChans[queries[0]],
+				runQueryStdin,
+			)
 		}
 	}()
 

shui.go

@@ -22,6 +22,10 @@ const (
 	MODE_READ                         // For running in 'read' mode.
 )
 
+const (
+	STDIN_QUERY_NAME = "stdin" // Named query value for reading stdin.
+)
+
 var (
 	ctx = context.Background() // Initialize context.
 )
@@ -37,8 +41,28 @@ func Run(config lib.Config, displayConfig lib.DisplayConfig) {
 	slog.Debug("Running with config", "config", config)
 	slog.Debug("Running with display config", "displayConfig", displayConfig)
 
+	// Define special query value when reading standard input.
+	if config.ReadStdin {
+		config.Queries = []string{STDIN_QUERY_NAME}
+	}
+
 	// Execute the specified mode.
 	switch {
+	case config.ReadStdin:
+		slog.Debug("Reading from standard input")
+
+		doneQueriesChan, pauseQueryChans = lib.Query(
+			lib.QUERY_MODE_STDIN,
+			-1,
+			config.Delay,
+			config.Queries,
+			config.Port,
+			config.History,
+			resultsReadyChan,
+		)
+
+		// Use labels that match the defined value for queries.
+		ctx = context.WithValue(ctx, "labels", config.Queries)
 	case config.Mode == int(MODE_PROFILE):
 		slog.Debug("Executing in profile mode")