Add --dump to discovery

cmd/livepeer/integration_test.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"crypto/sha256"
 	"encoding/json"
 	"fmt"
@@ -91,6 +92,7 @@ func TestRealBinaryOffchainFlow(t *testing.T) {
 	}
 
 	testBinaryMigrations(t, install)
+	t.Run("discovery dump", func(t *testing.T) { testDiscoveryDumpBinary(t, install) })
 
 	runner := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		require.NotEmpty(t, r.Header.Get("Livepeer-Session-Token"))
@@ -233,6 +235,161 @@ func testDiscoveryBinary(t *testing.T, install, orchestratorURL string) {
 	}, 5*time.Second, 20*time.Millisecond)
 }
 
+func testDiscoveryDumpBinary(t *testing.T, install string) {
+	t.Helper()
+	listen, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer listen.Close()
+	metrics, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer metrics.Close()
+	started := make(chan error, 1)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/inventory/discovery":
+			_, _ = io.WriteString(w, `[{"address":"https://orch.example","runners":[{"url":"https://orch.example/apps/runner/app","app":"test","observed_at":"2000-01-01T00:00:00Z"}]}]`)
+		case "/blocked/a/discovery", "/blocked/b/discovery":
+			_, _ = fmt.Fprintf(w, `[{"address":"https://orch.example","runners":[{"url":%q,"app":"test","observed_at":"2000-01-01T00:00:00Z","extension":%q}]}]`, "https://orch.example"+r.URL.Path, strings.Repeat("x", 900<<10))
+		case "/failed/discovery":
+			w.WriteHeader(http.StatusServiceUnavailable)
+		case "/cancel/discovery":
+			started <- nil
+			<-r.Context().Done()
+		default:
+			t.Errorf("unexpected request: %s", r.URL.Path)
+			w.WriteHeader(http.StatusNotFound)
+		}
+	}))
+	defer upstream.Close()
+
+	for _, entry := range [][]string{{"livepeer-discovery"}, {"livepeer", "discovery"}} {
+		t.Run(entry[0], func(t *testing.T) {
+			for _, tc := range []struct {
+				name       string
+				sources    []string
+				flags      []string
+				entries    int
+				exit       int
+				diagnostic string
+				noOutput   bool
+			}{
+				{name: "inventory", sources: []string{"inventory"}, entries: 2},
+				{name: "partial failure", sources: []string{"inventory", "failed"}, entries: 2, diagnostic: "discovery refresh failed"},
+				{name: "total failure", sources: []string{"failed"}, exit: 2, diagnostic: "all discovery sources failed"},
+				{name: "conflicting flags", flags: []string{"--print-config"}, exit: 2, diagnostic: "mutually exclusive", noOutput: true},
+				{name: "broken stdout", exit: 2, diagnostic: "cannot write discovery inventory", noOutput: true},
+				{name: "blocked stdout", sources: []string{"blocked/a", "blocked/b"}, exit: 2, diagnostic: "cannot write discovery inventory", noOutput: true},
+				{name: "reader closes after cancellation", sources: []string{"blocked/a"}, exit: 2, diagnostic: "cannot write discovery inventory", noOutput: true},
+				{name: "cancellation", sources: []string{"cancel"}, exit: 2, diagnostic: "context canceled", noOutput: true},
+			} {
+				t.Run(tc.name, func(t *testing.T) {
+					ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+					defer cancel()
+					args := append([]string{}, entry[1:]...)
+					args = append(args, "--config=", "--dump", "--listen", listen.Addr().String(), "--metrics-listen", metrics.Addr().String(),
+						"--grants", strings.TrimPrefix(upstream.URL, "http://"))
+					for _, source := range tc.sources {
+						args = append(args, "--nodes", upstream.URL+"/"+source)
+					}
+					args = append(args, tc.flags...)
+					cmd := exec.CommandContext(ctx, filepath.Join(install, entry[0]), args...)
+					var out, logs bytes.Buffer
+					cmd.Stdout, cmd.Stderr = &out, &logs
+					var ready chan error
+					if tc.name == "cancellation" {
+						ready = started
+					}
+					var stdoutReader, stderrReader, stderrWriter *os.File
+					if tc.name == "broken stdout" || tc.name == "blocked stdout" || tc.name == "reader closes after cancellation" {
+						reader, writer, err := os.Pipe()
+						require.NoError(t, err)
+						defer reader.Close()
+						defer writer.Close()
+						cmd.Stdout = writer
+						stdoutReader = reader
+						if tc.name == "broken stdout" {
+							require.NoError(t, reader.Close())
+						} else {
+							ready = make(chan error, 1)
+							go func() {
+								var first [1]byte
+								_, err := reader.Read(first[:])
+								ready <- err
+							}()
+						}
+					}
+					if tc.name == "reader closes after cancellation" {
+						stderrReader, stderrWriter, err = os.Pipe()
+						require.NoError(t, err)
+						defer stderrReader.Close()
+						defer stderrWriter.Close()
+						require.NoError(t, stderrWriter.SetWriteDeadline(time.Now().Add(100*time.Millisecond)))
+						_, err := stderrWriter.Write(make([]byte, 1<<20))
+						require.ErrorIs(t, err, os.ErrDeadlineExceeded)
+						require.NoError(t, stderrWriter.SetWriteDeadline(time.Time{}))
+						cmd.Stderr = stderrWriter
+					}
+					require.NoError(t, cmd.Start())
+					if stderrWriter != nil {
+						require.NoError(t, stderrWriter.Close())
+					}
+					done := make(chan error, 1)
+					go func() { done <- cmd.Wait() }()
+					if ready != nil {
+						select {
+						case err := <-ready:
+							require.NoError(t, err)
+							require.NoError(t, cmd.Process.Signal(syscall.SIGTERM))
+						case <-ctx.Done():
+							t.Fatal("dump did not become ready for cancellation")
+						}
+					}
+					if stderrReader != nil {
+						// Delay the final error report until cancellation has returned
+						// from Root and the stdout reader has closed.
+						for i := range 2 {
+							select {
+							case err := <-done:
+								t.Fatalf("dump exited before stderr was drained: %v", err)
+							case <-time.After(100 * time.Millisecond):
+							}
+							if i == 0 {
+								require.NoError(t, stdoutReader.Close())
+							}
+						}
+						_, err := io.Copy(&logs, stderrReader)
+						require.NoError(t, err)
+					}
+					err := <-done
+					require.NoError(t, ctx.Err(), "dump must finish before its process deadline")
+					if tc.exit == 0 {
+						require.NoError(t, err, logs.String())
+					} else {
+						var exit *exec.ExitError
+						require.ErrorAs(t, err, &exit, logs.String())
+						require.Equal(t, tc.exit, exit.ExitCode(), logs.String())
+					}
+					if tc.diagnostic == "" {
+						require.Empty(t, logs.String())
+					} else {
+						require.Contains(t, logs.String(), tc.diagnostic)
+					}
+					if tc.noOutput {
+						require.Empty(t, out.String())
+						return
+					}
+					var entries []json.RawMessage
+					require.NoError(t, json.Unmarshal(out.Bytes(), &entries), out.String())
+					require.NotNil(t, entries)
+					require.Len(t, entries, tc.entries)
+					require.True(t, strings.HasSuffix(out.String(), "\n"))
+					require.Equal(t, 1, strings.Count(out.String(), "\n"))
+				})
+			}
+		})
+	}
+}
+
 func checkPythonSDK(t *testing.T, root, serviceURL, runnerURL string) {
 	t.Helper()
 	python, export := test.PythonSDK(t, root)

cmd/livepeer/testdata/cli/discovery-help.golden

@@ -16,5 +16,6 @@ Flags:
       --max-failed-refreshes int    Consecutive failed refreshes before eviction (env: LIVEPEER_DISCOVERY_MAX_FAILED_REFRESHES) (default 6)
   -g, --grants strings              Private discovery IP[:port] or CIDR grants (env: LIVEPEER_DISCOVERY_GRANTS)
   -p, --print-config                Print TOML configuration
+      --dump                        Query discovery sources once, print JSON inventory, and exit
   -h, --help                        help for livepeer-discovery
   -v, --version                     version for livepeer-discovery

discovery/app.go

@@ -11,6 +11,7 @@ import (
 	"net/http"
 	"net/netip"
 	"net/url"
+	"os"
 	"os/signal"
 	"sync"
 	"syscall"
@@ -35,6 +36,7 @@ type Params struct {
 	MaxFailedRefreshes int                  `default:"6" min:"0" descr:"Consecutive failed refreshes before eviction"`
 	Grants             []string             `optional:"true" descr:"Private discovery IP[:port] or CIDR grants"`
 	PrintConfig        bool                 `optional:"true" boa:"noconfig,noenv" descr:"Print TOML configuration"`
+	Dump               bool                 `optional:"true" boa:"noconfig,noenv" descr:"Query discovery sources once, print JSON inventory, and exit"`
 }
 
 func Root(out, errOut io.Writer) *cobra.Command {
@@ -44,6 +46,9 @@ func Root(out, errOut io.Writer) *cobra.Command {
 		Params: params, RejectUnknown: true,
 		Args: cobra.NoArgs,
 		RunFuncE: func(p *Params, cmd *cobra.Command, _ []string) error {
+			if p.Dump && p.PrintConfig {
+				return boa.NewUserInputErrorf("--dump and --print-config are mutually exclusive")
+			}
 			if p.PrintConfig {
 				app, err := NewServer(*p, slog.Default())
 				if err != nil {
@@ -57,6 +62,19 @@ func Root(out, errOut io.Writer) *cobra.Command {
 					"RefreshInterval": &p.RefreshInterval, "MaxFailedRefreshes": &p.MaxFailedRefreshes,
 				})
 			}
+			if p.Dump {
+				ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
+				defer stop()
+				app, err := NewServer(*p, slog.New(slog.NewTextHandler(cmd.ErrOrStderr(), nil)))
+				if err != nil {
+					return boa.NewUserInputError(err)
+				}
+				inventoryOut := cmd.OutOrStdout()
+				if inventoryOut == os.Stdout {
+					inventoryOut = dumpStdout{ctx}
+				}
+				return app.dump(ctx, inventoryOut)
+			}
 			return Serve(cmd.Context(), *p, cmd.ErrOrStderr())
 		},
 	}, &params.ConfigFile).ToCobra()
@@ -68,6 +86,34 @@ func Root(out, errOut io.Writer) *cobra.Command {
 	return cmd
 }
 
+// dumpStdout lets a signal abort a blocked CLI write. The process exits after
+// dump returns, so this must only wrap process-owned stdout.
+type dumpStdout struct{ ctx context.Context }
+
+func (w dumpStdout) Write(data []byte) (int, error) {
+	type result struct {
+		n   int
+		err error
+	}
+	done := make(chan result, 1)
+	go func() {
+		brokenPipe := make(chan os.Signal, 1)
+		signal.Notify(brokenPipe, syscall.SIGPIPE)
+		defer signal.Stop(brokenPipe)
+		n, err := os.Stdout.Write(data)
+		done <- result{n, err}
+	}()
+	select {
+	case <-w.ctx.Done():
+		return 0, w.ctx.Err()
+	case written := <-done:
+		if err := w.ctx.Err(); err != nil {
+			return written.n, err
+		}
+		return written.n, written.err
+	}
+}
+
 func Serve(parent context.Context, p Params, logOut io.Writer) (result error) {
 	logger := slog.New(slog.NewTextHandler(logOut, nil))
 	app, err := NewServer(p, logger)