Make osq runner responsive to registration updates

cmd/launcher/launcher.go

@@ -373,7 +373,7 @@ func runLauncher(ctx context.Context, cancel func(), multiSlogger, systemMultiSl
 		osqueryruntime.WithAugeasLensFunction(augeas.InstallLenses),
 	)
 	runGroup.Add("osqueryRunner", osqueryRunner.Run, osqueryRunner.Interrupt)
-	k.SetInstanceQuerier(osqueryRunner)
+	k.SetInstanceRunner(osqueryRunner)
 
 	versionInfo := version.Version()
 	k.SystemSlogger().Log(ctx, slog.LevelInfo,

ee/agent/knapsack/knapsack.go

@@ -39,7 +39,7 @@ type knapsack struct {
 
 	slogger, systemSlogger *multislogger.MultiSlogger
 
-	querier types.InstanceQuerier
+	osqRunner types.OsqRunner
 
 	// This struct is a work in progress, and will be iteratively added to as needs arise.
 }
@@ -87,23 +87,31 @@ func (k *knapsack) AddSlogHandler(handler ...slog.Handler) {
 	k.systemSlogger.AddHandler(handler...)
 }
 
-// Osquery instance querier
-func (k *knapsack) SetInstanceQuerier(q types.InstanceQuerier) {
-	k.querier = q
+// Osquery instance runner
+func (k *knapsack) SetInstanceRunner(r types.OsqRunner) {
+	k.osqRunner = r
 }
 
 // RegistrationTracker interface methods
 func (k *knapsack) RegistrationIDs() []string {
 	return []string{types.DefaultRegistrationID}
 }
 
+func (k *knapsack) SetRegistrationIDs(registrationIDs []string) error {
+	if k.osqRunner == nil {
+		return nil
+	}
+
+	return k.osqRunner.UpdateRegistrationIDs(registrationIDs)
+}
+
 // InstanceStatuses returns the current status of each osquery instance.
 // It performs a healthcheck against each existing instance.
 func (k *knapsack) InstanceStatuses() map[string]types.InstanceStatus {
-	if k.querier == nil {
+	if k.osqRunner == nil {
 		return nil
 	}
-	return k.querier.InstanceStatuses()
+	return k.osqRunner.InstanceStatuses()
 }
 
 // BboltDB interface methods

ee/agent/types/knapsack.go

@@ -11,7 +11,7 @@ type Knapsack interface {
 	Slogger
 	RegistrationTracker
 	InstanceQuerier
-	SetInstanceQuerier(q InstanceQuerier)
+	SetInstanceRunner(r OsqRunner)
 	// LatestOsquerydPath finds the path to the latest osqueryd binary, after accounting for updates.
 	LatestOsquerydPath(ctx context.Context) string
 	// ReadEnrollSecret returns the enroll secret value, checking in various locations.

ee/agent/types/registration.go

@@ -9,4 +9,5 @@ const (
 // data may be provided by e.g. a control server subsystem.
 type RegistrationTracker interface {
 	RegistrationIDs() []string
+	SetRegistrationIDs(registrationIDs []string) error
 }

ee/agent/types/runner.go

@@ -0,0 +1,13 @@
+package types
+
+type (
+	// RegistrationChangeHandler is implemented by pkg/osquery/runtime/runner.go
+	RegistrationChangeHandler interface {
+		UpdateRegistrationIDs(registrationIDs []string) error
+	}
+
+	OsqRunner interface {
+		RegistrationChangeHandler
+		InstanceQuerier
+	}
+)

pkg/osquery/runtime/runner.go

@@ -5,7 +5,9 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"slices"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/kolide/launcher/ee/agent/flags/keys"
@@ -26,8 +28,9 @@ type Runner struct {
 	knapsack        types.Knapsack
 	serviceClient   service.KolideService   // shared service client for communication between osquery instance and Kolide SaaS
 	opts            []OsqueryInstanceOption // global options applying to all osquery instances
-	shutdown        chan struct{}
-	interrupted     bool
+	shutdown        chan struct{}           // buffered shutdown channel for to enable shutting down to restart or exit
+	rerunRequired   atomic.Bool
+	interrupted     atomic.Bool
 }
 
 func New(k types.Knapsack, serviceClient service.KolideService, opts ...OsqueryInstanceOption) *Runner {
@@ -37,8 +40,9 @@ func New(k types.Knapsack, serviceClient service.KolideService, opts ...OsqueryI
 		slogger:         k.Slogger().With("component", "osquery_runner"),
 		knapsack:        k,
 		serviceClient:   serviceClient,
-		shutdown:        make(chan struct{}),
-		opts:            opts,
+		// the buffer length is arbitrarily set at 100, this number just needs to be higher than the total possible instances
+		shutdown: make(chan struct{}, 100),
+		opts:     opts,
 	}
 
 	k.RegisterChangeObserver(runner,
@@ -49,6 +53,31 @@ func New(k types.Knapsack, serviceClient service.KolideService, opts ...OsqueryI
 }
 
 func (r *Runner) Run() error {
+	for {
+		// if our instances ever exit unexpectedly, return immediately
+		if err := r.runRegisteredInstances(); err != nil {
+			return err
+		}
+
+		// if we're in a state that required re-running all registered instances,
+		// reset the field and do that
+		if r.rerunRequired.Load() {
+			r.rerunRequired.Store(false)
+			continue
+		}
+
+		// otherwise, exit cleanly
+		return nil
+	}
+}
+
+func (r *Runner) runRegisteredInstances() error {
+	// clear the internal instances to add back in fresh as we runInstance,
+	// this prevents old instances from sticking around if a registrationID is ever removed
+	r.instanceLock.Lock()
+	r.instances = make(map[string]*OsqueryInstance)
+	r.instanceLock.Unlock()
+
 	// Create a group to track the workers running each instance
 	wg, ctx := errgroup.WithContext(context.Background())
 
@@ -186,6 +215,13 @@ func (r *Runner) Query(query string) ([]map[string]string, error) {
 }
 
 func (r *Runner) Interrupt(_ error) {
+	if r.interrupted.Load() {
+		// Already shut down, nothing else to do
+		return
+	}
+
+	r.interrupted.Store(true)
+
 	if err := r.Shutdown(); err != nil {
 		r.slogger.Log(context.TODO(), slog.LevelWarn,
 			"could not shut down runner on interrupt",
@@ -197,13 +233,12 @@ func (r *Runner) Interrupt(_ error) {
 // Shutdown instructs the runner to permanently stop the running instance (no
 // restart will be attempted).
 func (r *Runner) Shutdown() error {
-	if r.interrupted {
-		// Already shut down, nothing else to do
-		return nil
+	// ensure one shutdown is sent for each instance to read
+	r.instanceLock.Lock()
+	for range r.instances {
+		r.shutdown <- struct{}{}
 	}
-
-	r.interrupted = true
-	close(r.shutdown)
+	r.instanceLock.Unlock()
 
 	if err := r.triggerShutdownForInstances(); err != nil {
 		return fmt.Errorf("triggering shutdown for instances during runner shutdown: %w", err)
@@ -333,3 +368,42 @@ func (r *Runner) InstanceStatuses() map[string]types.InstanceStatus {
 
 	return instanceStatuses
 }
+
+// UpdateRegistrationIDs detects any changes between the new and stored registration IDs,
+// and resets the runner instances for the new registrationIDs if required
+func (r *Runner) UpdateRegistrationIDs(newRegistrationIDs []string) error {
+	slices.Sort(newRegistrationIDs)
+	existingRegistrationIDs := r.registrationIds
+	slices.Sort(existingRegistrationIDs)
+
+	if slices.Equal(newRegistrationIDs, existingRegistrationIDs) {
+		r.slogger.Log(context.TODO(), slog.LevelDebug,
+			"skipping runner restarts for updated registration IDs, no changes detected",
+		)
+
+		return nil
+	}
+
+	r.slogger.Log(context.TODO(), slog.LevelDebug,
+		"detected changes to registrationIDs, will restart runner instances",
+		"previous_registration_ids", existingRegistrationIDs,
+		"new_registration_ids", newRegistrationIDs,
+	)
+
+	// we know there are changes, safe to update the internal registrationIDs now
+	r.registrationIds = newRegistrationIDs
+	// mark rerun as required so that we can safely shutdown all workers and have the changes
+	// picked back up from within the main Run function
+	r.rerunRequired.Store(true)
+
+	if err := r.Shutdown(); err != nil {
+		r.slogger.Log(context.TODO(), slog.LevelWarn,
+			"could not shut down runner instances for restart after registration changes",
+			"err", err,
+		)
+
+		return err
+	}
+
+	return nil
+}