feat(plugins.nvidia_smi): Implement Probe()

agent/agent.go

@@ -378,6 +378,12 @@ func (a *Agent) startInputs(
 
 			return nil, fmt.Errorf("starting input %s: %w", input.LogName(), err)
 		}
+		if err := input.Probe(); err != nil {
+			// Probe failures are not fatal to Telegraf itself, so simply skip
+			// adding the input.
+			log.Printf("I! [agent] Failed to probe %s, shutting down plugin: %s", input.LogName(), err)
+			continue
+		}
 		unit.inputs = append(unit.inputs, input)
 	}
 

agent/agent_test.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -257,3 +258,53 @@ func collect(ctx context.Context, a *Agent, wait time.Duration) ([]telegraf.Metr
 	}
 	return received, nil
 }
+
+type mockProbingInput struct {
+	probeReturn error
+}
+
+func (m *mockProbingInput) SampleConfig() string {
+	return ""
+}
+
+func (m *mockProbingInput) Gather(_ telegraf.Accumulator) error {
+	return nil
+}
+
+func (m *mockProbingInput) Probe() error {
+	return m.probeReturn
+}
+
+func TestAgentstartInputsProbing(t *testing.T) {
+	for _, tt := range []struct {
+		name                string
+		probeReturnError    error
+		expectedInputLength int
+	}{
+		{
+			name:                "probe failed",
+			probeReturnError:    errors.New("probe failure"),
+			expectedInputLength: 0,
+		},
+		{
+			name:                "probe succeeded",
+			expectedInputLength: 1,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			a := NewAgent(config.NewConfig())
+			require.Empty(t, a.Config.Outputs)
+
+			inputs := []*models.RunningInput{
+				models.NewRunningInput(&mockProbingInput{
+					probeReturn: tt.probeReturnError,
+				}, &models.InputConfig{
+					StartupErrorBehavior: "probe",
+				}),
+			}
+			inputUnit, err := a.startInputs(make(chan<- telegraf.Metric), inputs)
+			require.NoError(t, err)
+			require.Len(t, inputUnit.inputs, tt.expectedInputLength)
+		})
+	}
+}

models/running_input.go

@@ -112,7 +112,7 @@ func (r *RunningInput) LogName() string {
 
 func (r *RunningInput) Init() error {
 	switch r.Config.StartupErrorBehavior {
-	case "", "error", "retry", "ignore":
+	case "", "error", "retry", "ignore", "probe":
 	default:
 		return fmt.Errorf("invalid 'startup_error_behavior' setting %q", r.Config.StartupErrorBehavior)
 	}
@@ -161,7 +161,7 @@ func (r *RunningInput) Start(acc telegraf.Accumulator) error {
 		}
 		r.log.Infof("Startup failed: %v; retrying...", err)
 		return nil
-	case "ignore":
+	case "ignore", "probe":
 		return &internal.FatalError{Err: serr}
 	default:
 		r.log.Errorf("Invalid 'startup_error_behavior' setting %q", r.Config.StartupErrorBehavior)
@@ -170,6 +170,16 @@ func (r *RunningInput) Start(acc telegraf.Accumulator) error {
 	return err
 }
 
+func (r *RunningInput) Probe() error {
+	p, ok := r.Input.(telegraf.ProbePlugin)
+	if !ok || r.Config.StartupErrorBehavior != "probe" {
+		r.log.Debug("Not probing plugin")
+		return nil
+	}
+	r.log.Debug("Probing plugin")
+	return p.Probe()
+}
+
 func (r *RunningInput) Stop() {
 	if plugin, ok := r.Input.(telegraf.ServiceInput); ok {
 		plugin.Stop()

models/running_input_test.go

@@ -1,9 +1,11 @@
 package models
 
 import (
+	"errors"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/influxdata/telegraf"
@@ -487,6 +489,71 @@ func TestRunningInputMakeMetricWithGatherEndTimeSource(t *testing.T) {
 	require.Equal(t, expected, actual)
 }
 
+func TestRunningInputProbing(t *testing.T) {
+	probeErr := errors.New("probing error")
+	for _, tt := range []struct {
+		name                 string
+		input                telegraf.Input
+		startupErrorBehavior string
+		expectedError        bool
+	}{
+		{
+			name:                 "non-probing plugin with probe value set",
+			input:                &mockInput{},
+			startupErrorBehavior: "probe",
+			expectedError:        false,
+		},
+		{
+			name:                 "non-probing plugin with probe value not set",
+			input:                &mockInput{},
+			startupErrorBehavior: "ignore",
+			expectedError:        false,
+		},
+		{
+			name:                 "probing plugin with probe value set",
+			input:                &mockProbingInput{probeErr},
+			startupErrorBehavior: "probe",
+			expectedError:        true,
+		},
+		{
+			name:                 "probing plugin with probe value not set",
+			input:                &mockProbingInput{probeErr},
+			startupErrorBehavior: "ignore",
+			expectedError:        false,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			ri := NewRunningInput(tt.input, &InputConfig{
+				Name:                 "TestRunningInput",
+				StartupErrorBehavior: tt.startupErrorBehavior,
+			})
+			ri.log = testutil.Logger{}
+			err := ri.Probe()
+			if tt.expectedError {
+				assert.Error(t, err)
+			} else {
+				assert.NoError(t, err)
+			}
+		})
+	}
+}
+
+type mockProbingInput struct {
+	probeReturn error
+}
+
+func (m *mockProbingInput) SampleConfig() string {
+	return ""
+}
+
+func (m *mockProbingInput) Gather(_ telegraf.Accumulator) error {
+	return nil
+}
+
+func (m *mockProbingInput) Probe() error {
+	return m.probeReturn
+}
+
 type mockInput struct{}
 
 func (t *mockInput) SampleConfig() string {

plugin.go

@@ -56,3 +56,9 @@ type StatefulPlugin interface {
 	// initialization (after Init() function).
 	SetState(state interface{}) error
 }
+
+// ProbePlugin is an interface that all input/output plugins need to
+// implement in order to support the `probe` value of `startup_error_behavior`.
+type ProbePlugin interface {
+	Probe() error
+}

plugins/inputs/nvidia_smi/nvidia_smi.go

@@ -31,8 +31,9 @@ type NvidiaSMI struct {
 	Timeout config.Duration `toml:"timeout"`
 	Log     telegraf.Logger `toml:"-"`
 
-	ignorePlugin bool
-	once         sync.Once
+	nvidiaSMIArgs []string
+	ignorePlugin  bool
+	once          sync.Once
 }
 
 func (*NvidiaSMI) SampleConfig() string {
@@ -53,14 +54,23 @@ func (smi *NvidiaSMI) Start(telegraf.Accumulator) error {
 
 func (smi *NvidiaSMI) Stop() {}
 
+func (smi *NvidiaSMI) Probe() error {
+	// Construct and execute metrics query
+	_, err := internal.CombinedOutputTimeout(exec.Command(smi.BinPath, smi.nvidiaSMIArgs...), time.Duration(smi.Timeout))
+	if err != nil {
+		return fmt.Errorf("calling %q failed: %w", smi.BinPath, err)
+	}
+	return nil
+}
+
 // Gather implements the telegraf interface
 func (smi *NvidiaSMI) Gather(acc telegraf.Accumulator) error {
 	if smi.ignorePlugin {
 		return nil
 	}
 
 	// Construct and execute metrics query
-	data, err := internal.CombinedOutputTimeout(exec.Command(smi.BinPath, "-q", "-x"), time.Duration(smi.Timeout))
+	data, err := internal.CombinedOutputTimeout(exec.Command(smi.BinPath, smi.nvidiaSMIArgs...), time.Duration(smi.Timeout))
 	if err != nil {
 		return fmt.Errorf("calling %q failed: %w", smi.BinPath, err)
 	}
@@ -119,8 +129,9 @@ func (smi *NvidiaSMI) parse(acc telegraf.Accumulator, data []byte) error {
 func init() {
 	inputs.Add("nvidia_smi", func() telegraf.Input {
 		return &NvidiaSMI{
-			BinPath: "/usr/bin/nvidia-smi",
-			Timeout: config.Duration(5 * time.Second),
+			BinPath:       "/usr/bin/nvidia-smi",
+			Timeout:       config.Duration(5 * time.Second),
+			nvidiaSMIArgs: []string{"-q", "-x"},
 		}
 	})
 }

plugins/inputs/nvidia_smi/nvidia_smi_test.go

@@ -4,16 +4,66 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"runtime"
 	"testing"
 	"time"
 
 	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/config"
 	"github.com/influxdata/telegraf/internal"
 	"github.com/influxdata/telegraf/models"
 	"github.com/influxdata/telegraf/testutil"
 	"github.com/stretchr/testify/require"
 )
 
+func TestProbe(t *testing.T) {
+	var binPath string
+	var nvidiaSMIArgsPrefix []string
+	if runtime.GOOS == "windows" {
+		binPath = `C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe`
+		nvidiaSMIArgsPrefix = []string{"-Command"}
+	} else {
+		binPath = "/bin/bash"
+		nvidiaSMIArgsPrefix = []string{"-c"}
+	}
+
+	for _, tt := range []struct {
+		name        string
+		args        string
+		expectError bool
+	}{
+		{
+			name:        "probe success",
+			args:        "exit 0",
+			expectError: false,
+		},
+		{
+			name:        "probe error",
+			args:        "exit 1",
+			expectError: true,
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			plugin := &NvidiaSMI{
+				BinPath:       binPath,
+				nvidiaSMIArgs: append(nvidiaSMIArgsPrefix, tt.args),
+				Log:           &testutil.Logger{},
+				Timeout:       config.Duration(5 * time.Second),
+			}
+			model := models.NewRunningInput(plugin, &models.InputConfig{
+				Name:                 "nvidia_smi",
+				StartupErrorBehavior: "probe",
+			})
+			err := model.Probe()
+			if tt.expectError {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
 func TestErrorBehaviorDefault(t *testing.T) {
 	// make sure we can't find nvidia-smi in $PATH somewhere
 	os.Unsetenv("PATH")