make provider reconnectable

providers/coordinator.go

@@ -9,7 +9,6 @@ import (
  "os/exec"
  "strconv"
  "sync"
- "time"
 
  "github.com/cockroachdb/errors"
  "github.com/hashicorp/go-hclog"
@@ -310,51 +309,51 @@ func (c *coordinator) unsafeStartProvider(id string, update UpdateProvidersConfi
  }
  }
 
- pluginCmd := exec.Command(provider.binPath(), []string{"run_as_plugin", "--log-level", zerolog.GlobalLevel().String()}...)
-
- addColorConfig(pluginCmd)
-
- pluginLogger := &hclogger{Logger: log.Logger}
- pluginLogger.SetLevel(hclog.Warn)
- client := plugin.NewClient(&plugin.ClientConfig{
- HandshakeConfig: pp.Handshake,
- Plugins: pp.PluginMap,
- Cmd: pluginCmd,
- AllowedProtocols: []plugin.Protocol{
- plugin.ProtocolNetRPC, plugin.ProtocolGRPC,
- },
- Logger: pluginLogger,
- Stderr: os.Stderr,
- })
-
- // Connect via RPC
- rpcClient, err := client.Client()
- if err != nil {
- client.Kill()
- return nil, errors.Wrap(err, "failed to initialize plugin client")
- }
+ connectFunc := func() (pp.ProviderPlugin, *plugin.Client, error) {
+ pluginCmd := exec.Command(provider.binPath(), []string{"run_as_plugin", "--log-level", zerolog.GlobalLevel().String()}...)
+
+ addColorConfig(pluginCmd)
+
+ pluginLogger := &hclogger{Logger: log.Logger}
+ pluginLogger.SetLevel(hclog.Warn)
+ client := plugin.NewClient(&plugin.ClientConfig{
+ HandshakeConfig: pp.Handshake,
+ Plugins: pp.PluginMap,
+ Cmd: pluginCmd,
+ AllowedProtocols: []plugin.Protocol{
+ plugin.ProtocolNetRPC, plugin.ProtocolGRPC,
+ },
+ Logger: pluginLogger,
+ Stderr: os.Stderr,
+ })
+
+ // Connect via RPC
+ rpcClient, err := client.Client()
+ if err != nil {
+ client.Kill()
+ return nil, nil, errors.Wrap(err, "failed to initialize plugin client")
+ }
 
- // Request the plugin
- pluginName := "provider"
- raw, err := rpcClient.Dispense(pluginName)
- if err != nil {
- client.Kill()
- return nil, errors.Wrap(err, "failed to call "+pluginName+" plugin")
+ // Request the plugin
+ pluginName := "provider"
+ raw, err := rpcClient.Dispense(pluginName)
+ if err != nil {
+ client.Kill()
+ return nil, nil, errors.Wrap(err, "failed to call "+pluginName+" plugin")
+ }
+
+ return raw.(pp.ProviderPlugin), client, nil
  }
 
- res := &RunningProvider{
- Name: provider.Name,
- ID: provider.ID,
- Plugin: raw.(pp.ProviderPlugin),
- Client: client,
- Schema: provider.Schema,
- interval: 2 * time.Second,
- gracePeriod: 3 * time.Second,
+ plug, client, err := connectFunc()
+ if err != nil {
+ return nil, err
  }
 
  c.schema.Add(provider.ID, provider.Schema)
 
- if err := res.heartbeat(); err != nil {
+ res, err := SupervisedRunningProivder(provider.Name, provider.ID, plug, client, provider.Schema, connectFunc)
+ if err != nil {
  return nil, err
  }
  c.runningByID[res.ID] = res
@@ -388,7 +387,7 @@ func (c *coordinator) Shutdown() {
  if err := provider.Shutdown(); err != nil {
  log.Warn().Err(err).Str("provider", provider.Name).Msg("failed to shut down provider")
  }
- provider.Client.Kill()
+ provider.KillClient()
  }
  c.runningByID = map[string]*RunningProvider{}
  c.runtimes = map[string]*Runtime{}

providers/coordinator_test.go

@@ -7,7 +7,6 @@ import (
  "fmt"
  "testing"
 
- "github.com/hashicorp/go-plugin"
  "github.com/stretchr/testify/assert"
  "go.mondoo.com/cnquery/v11/providers-sdk/v1/inventory"
  pp "go.mondoo.com/cnquery/v11/providers-sdk/v1/plugin"
@@ -32,7 +31,6 @@ func TestShutdown(t *testing.T) {
  c.runningByID[id] = &RunningProvider{
  ID: id,
  Plugin: mockPlugin,
- Client: &plugin.Client{},
  }
  }
 

providers/providers_test.go

@@ -4,6 +4,7 @@
 package providers
 
 import (
+ "context"
  "syscall"
  "testing"
  "time"
@@ -50,7 +51,9 @@ func TestProviderShutdown(t *testing.T) {
  interval: 500 * time.Millisecond,
  gracePeriod: 500 * time.Millisecond,
  }
- err := s.heartbeat()
+ hbtCtx, hbtCancel := context.WithCancel(context.Background())
+ s.hbCancelFunc = hbtCancel
+ err := s.heartbeat(hbtCtx, hbtCancel)
  require.NoError(t, err)
  require.False(t, s.isCloseOrShutdown())
  // the shutdown here takes 10 seconds, whereas the heartbeat interval is every second.