feat: add ability to launch targeted dskit modules in the grafana server CLI command (#74188)

* feat: add ability to launch targeted dskit modules in the grafana server CLI command

This commit adds a ModuleServer and ModuleRunner suitable for launching dskit services and updates the server cli command to use this instead of the full Server. The default behavior is unchanged and will launch the full Grafana server. Individual services are targeted by setting target=comma,seperated,list in the config file.

* require dev mode to target dskit modules

* remove unused type

* replace setting.CommandLineArgs w/setting.Cfg; the caller can deal with calling setting.NewCfg

* Update pkg/server/module_server.go

Co-authored-by: Serge Zaitsev <serge.zaitsev@grafana.com>

---------

Co-authored-by: Serge Zaitsev <serge.zaitsev@grafana.com>

go.sum

@@ -1763,8 +1763,6 @@ github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWm
 github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/gotestyourself/gotestyourself v1.3.0/go.mod h1:zZKM6oeNM8k+FRljX1mnzVYeS8wiGgQyvST1/GafPbY=
 github.com/gotestyourself/gotestyourself v2.2.0+incompatible/go.mod h1:zZKM6oeNM8k+FRljX1mnzVYeS8wiGgQyvST1/GafPbY=
-github.com/grafana/alerting v0.0.0-20230825102000-717e20092271 h1:Nl7/Da/OLKrB5MpSKzSqQsshOIXmXjzcUBZ83ge1sKY=
-github.com/grafana/alerting v0.0.0-20230825102000-717e20092271/go.mod h1:gyUqgDT+v6gARVCpbfi8bb/WiGNELNJiq6hGKadnIxc=
 github.com/grafana/alerting v0.0.0-20230831092459-f7dba8ede9b0 h1:iDCcF7YzzVWQDtgUp3pS5caAAfI3lk7RZ0BI5DMrero=
 github.com/grafana/alerting v0.0.0-20230831092459-f7dba8ede9b0/go.mod h1:gyUqgDT+v6gARVCpbfi8bb/WiGNELNJiq6hGKadnIxc=
 github.com/grafana/codejen v0.0.3 h1:tAWxoTUuhgmEqxJPOLtJoxlPBbMULFwKFOcRsPRPXDw=

pkg/cmd/grafana-cli/commands/commands.go

@@ -50,12 +50,17 @@ func runDbCommand(command func(commandLine utils.CommandLine, sqlStore db.DB) er
 
 func initializeRunner(cmd *utils.ContextCommandLine) (server.Runner, error) {
 	configOptions := strings.Split(cmd.String("configOverrides"), " ")
-	runner, err := server.InitializeForCLI(setting.CommandLineArgs{
+	cfg, err := setting.NewCfgFromArgs(setting.CommandLineArgs{
 		Config:   cmd.ConfigFile(),
 		HomePath: cmd.HomePath(),
 		// tailing arguments have precedence over the options string
 		Args: append(configOptions, cmd.Args().Slice()...),
 	})
+	if err != nil {
+		return server.Runner{}, err
+	}
+
+	runner, err := server.InitializeForCLI(cfg)
 	if err != nil {
 		return server.Runner{}, fmt.Errorf("%v: %w", "failed to initialize runner", err)
 	}

pkg/cmd/grafana-server/commands/cli.go

@@ -93,14 +93,18 @@ func RunServer(opts ServerOptions) error {
 	checkPrivileges()
 
 	configOptions := strings.Split(ConfigOverrides, " ")
+	cfg, err := setting.NewCfgFromArgs(setting.CommandLineArgs{
+		Config:   ConfigFile,
+		HomePath: HomePath,
+		// tailing arguments have precedence over the options string
+		Args: append(configOptions, opts.Context.Args().Slice()...),
+	})
+	if err != nil {
+		return err
+	}
 
-	s, err := server.Initialize(
-		setting.CommandLineArgs{
-			Config:   ConfigFile,
-			HomePath: HomePath,
-			// tailing arguments have precedence over the options string
-			Args: append(configOptions, opts.Context.Args().Slice()...),
-		},
+	s, err := server.InitializeModuleServer(
+		cfg,
 		server.Options{
 			PidFile:     PidFile,
 			Version:     opts.Version,
@@ -114,9 +118,7 @@ func RunServer(opts ServerOptions) error {
 	}
 
 	ctx := context.Background()
-
 	go listenToSystemSignals(ctx, s)
-
 	return s.Run()
 }
 
@@ -130,7 +132,7 @@ func validPackaging(packaging string) string {
 	return "unknown"
 }
 
-func listenToSystemSignals(ctx context.Context, s *server.Server) {
+func listenToSystemSignals(ctx context.Context, s *server.ModuleServer) {
 	signalChan := make(chan os.Signal, 1)
 	sighupChan := make(chan os.Signal, 1)
 

pkg/server/doc.go

@@ -0,0 +1,20 @@
+// Server defines the main entrypoints to Grafana and the Grafana CLI, as well
+// as test environments. OSS and Enterprise-specific build tags are used in this
+// package to control wire dependencies for each build.
+package server
+
+// Notes about wiresets:
+//
+// wire.go contains wire sets used by both OSS and Enterprise builds. These are
+// generally base wiresets imported by the OSS- or Enterprise-specific sets.
+//
+// wireexts_oss.go contains the "extensions" wiresets, used only by OSS builds.
+// wireexts_enterprise.go contains wiresets used only by Enterprise builds. This
+// file is located in the grafana-enterprise repo.
+//
+// NOTE WELL: The extensions sets can import wiresets from wire.go, but sets in
+// wire.go cannot include a build-specific wireset. The extension set must be built in wire.go.
+//
+// We use go build tags during build to configure which wiresets are used in a
+// given build. We do not commit generated wire sets (wire_gen.go) into the
+// repo.

pkg/server/module_runner.go

@@ -0,0 +1,25 @@
+package server
+
+import (
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+// ModuleRunner is a simplified version of Runner that is used in the grafana
+// server target command. It has a minimal set of dependencies required to
+// launch background/dskit services.
+type ModuleRunner struct {
+	Cfg              *setting.Cfg
+	SettingsProvider setting.Provider
+	Features         featuremgmt.FeatureToggles
+}
+
+func NewModuleRunner(cfg *setting.Cfg, settingsProvider setting.Provider,
+	features featuremgmt.FeatureToggles,
+) ModuleRunner {
+	return ModuleRunner{
+		Cfg:              cfg,
+		SettingsProvider: settingsProvider,
+		Features:         features,
+	}
+}

pkg/server/module_server.go

@@ -0,0 +1,205 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"os"
+	"path"
+	"path/filepath"
+	"strconv"
+	"sync"
+
+	"github.com/grafana/dskit/services"
+
+	"github.com/grafana/grafana/pkg/api"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/modules"
+	grafanaapiserver "github.com/grafana/grafana/pkg/services/grafana-apiserver"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+// NewModule returns an instance of a ModuleServer, responsible for managing
+// dskit modules (services).
+func NewModule(opts Options, apiOpts api.ServerOptions, cfg *setting.Cfg) (*ModuleServer, error) {
+	s, err := newModuleServer(opts, apiOpts, cfg)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := s.init(); err != nil {
+		return nil, err
+	}
+
+	return s, nil
+}
+
+func newModuleServer(opts Options, apiOpts api.ServerOptions, cfg *setting.Cfg) (*ModuleServer, error) {
+	rootCtx, shutdownFn := context.WithCancel(context.Background())
+
+	s := &ModuleServer{
+		opts:             opts,
+		apiOpts:          apiOpts,
+		context:          rootCtx,
+		shutdownFn:       shutdownFn,
+		shutdownFinished: make(chan struct{}),
+		log:              log.New("base-server"),
+		cfg:              cfg,
+		pidFile:          opts.PidFile,
+		version:          opts.Version,
+		commit:           opts.Commit,
+		buildBranch:      opts.BuildBranch,
+	}
+
+	return s, nil
+}
+
+// ModuleServer is responsible for managing the lifecycle of dskit services. The
+// ModuleServer has the minimal set of dependencies to launch dskit services,
+// but it can be used to launch the entire Grafana server.
+type ModuleServer struct {
+	opts    Options
+	apiOpts api.ServerOptions
+
+	context          context.Context
+	shutdownFn       context.CancelFunc
+	log              log.Logger
+	cfg              *setting.Cfg
+	shutdownOnce     sync.Once
+	shutdownFinished chan struct{}
+	isInitialized    bool
+	mtx              sync.Mutex
+
+	pidFile     string
+	version     string
+	commit      string
+	buildBranch string
+}
+
+// init initializes the server and its services.
+func (s *ModuleServer) init() error {
+	s.mtx.Lock()
+	defer s.mtx.Unlock()
+
+	if s.isInitialized {
+		return nil
+	}
+	s.isInitialized = true
+
+	if err := s.writePIDFile(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+// Run initializes and starts services. This will block until all services have
+// exited. To initiate shutdown, call the Shutdown method in another goroutine.
+func (s *ModuleServer) Run() error {
+	defer close(s.shutdownFinished)
+
+	if err := s.init(); err != nil {
+		return err
+	}
+
+	s.notifySystemd("READY=1")
+	s.log.Debug("Waiting on services...")
+
+	// Only allow individual dskit modules to run in dev mode.
+	if s.cfg.Env != "dev" {
+		if len(s.cfg.Target) > 1 || s.cfg.Target[0] != "all" {
+			s.log.Error("dskit module targeting is only supported in dev mode. Falling back to 'all'")
+			s.cfg.Target = []string{"all"}
+		}
+	}
+
+	m := modules.New(s.cfg.Target)
+
+	m.RegisterModule(modules.Core, func() (services.Service, error) {
+		return NewService(s.cfg, s.opts, s.apiOpts)
+	})
+
+	m.RegisterModule(modules.GrafanaAPIServer, func() (services.Service, error) {
+		return grafanaapiserver.New(path.Join(s.cfg.DataPath, "k8s"))
+	})
+
+	m.RegisterModule(modules.All, nil)
+
+	return m.Run(s.context)
+}
+
+// Shutdown initiates Grafana graceful shutdown. This shuts down all
+// running background services. Since Run blocks Shutdown supposed to
+// be run from a separate goroutine.
+func (s *ModuleServer) Shutdown(ctx context.Context, reason string) error {
+	var err error
+	s.shutdownOnce.Do(func() {
+		s.log.Info("Shutdown started", "reason", reason)
+		// Call cancel func to stop background services.
+		s.shutdownFn()
+		// Wait for server to shut down
+		select {
+		case <-s.shutdownFinished:
+			s.log.Debug("Finished waiting for server to shut down")
+		case <-ctx.Done():
+			s.log.Warn("Timed out while waiting for server to shut down")
+			err = fmt.Errorf("timeout waiting for shutdown")
+		}
+	})
+
+	return err
+}
+
+// writePIDFile retrieves the current process ID and writes it to file.
+func (s *ModuleServer) writePIDFile() error {
+	if s.pidFile == "" {
+		return nil
+	}
+
+	// Ensure the required directory structure exists.
+	err := os.MkdirAll(filepath.Dir(s.pidFile), 0700)
+	if err != nil {
+		s.log.Error("Failed to verify pid directory", "error", err)
+		return fmt.Errorf("failed to verify pid directory: %s", err)
+	}
+
+	// Retrieve the PID and write it to file.
+	pid := strconv.Itoa(os.Getpid())
+	if err := os.WriteFile(s.pidFile, []byte(pid), 0644); err != nil {
+		s.log.Error("Failed to write pidfile", "error", err)
+		return fmt.Errorf("failed to write pidfile: %s", err)
+	}
+
+	s.log.Info("Writing PID file", "path", s.pidFile, "pid", pid)
+	return nil
+}
+
+// notifySystemd sends state notifications to systemd.
+func (s *ModuleServer) notifySystemd(state string) {
+	notifySocket := os.Getenv("NOTIFY_SOCKET")
+	if notifySocket == "" {
+		s.log.Debug(
+			"NOTIFY_SOCKET environment variable empty or unset, can't send systemd notification")
+		return
+	}
+
+	socketAddr := &net.UnixAddr{
+		Name: notifySocket,
+		Net:  "unixgram",
+	}
+	conn, err := net.DialUnix(socketAddr.Net, nil, socketAddr)
+	if err != nil {
+		s.log.Warn("Failed to connect to systemd", "err", err, "socket", notifySocket)
+		return
+	}
+	defer func() {
+		if err := conn.Close(); err != nil {
+			s.log.Warn("Failed to close connection", "err", err)
+		}
+	}()
+
+	_, err = conn.Write([]byte(state))
+	if err != nil {
+		s.log.Warn("Failed to write notification to systemd", "err", err)
+	}
+}

pkg/server/server.go

@@ -78,7 +78,9 @@ func newServer(opts Options, cfg *setting.Cfg, httpServer *api.HTTPServer, roleR
 	return s, nil
 }
 
-// Server is responsible for managing the lifecycle of services.
+// Server is responsible for managing the lifecycle of services. This is the
+// core Server implementation which starts the entire Grafana server. Use
+// ModuleServer to launch specific modules.
 type Server struct {
 	context          context.Context
 	shutdownFn       context.CancelFunc

pkg/server/server_test.go

@@ -90,30 +90,3 @@ func TestServer_Shutdown(t *testing.T) {
 	err = <-ch
 	require.NoError(t, err)
 }
-
-type MockModuleService struct {
-	initFunc     func(context.Context) error
-	runFunc      func(context.Context) error
-	shutdownFunc func(context.Context) error
-}
-
-func (m *MockModuleService) Init(ctx context.Context) error {
-	if m.initFunc != nil {
-		return m.initFunc(ctx)
-	}
-	return nil
-}
-
-func (m *MockModuleService) Run(ctx context.Context) error {
-	if m.runFunc != nil {
-		return m.runFunc(ctx)
-	}
-	return nil
-}
-
-func (m *MockModuleService) Shutdown(ctx context.Context) error {
-	if m.shutdownFunc != nil {
-		return m.shutdownFunc(ctx)
-	}
-	return nil
-}

pkg/server/service.go

@@ -11,23 +11,24 @@ import (
 
 type coreService struct {
 	*services.BasicService
-	cla     setting.CommandLineArgs
+	cfg     *setting.Cfg
 	opts    Options
 	apiOpts api.ServerOptions
 	server  *Server
 }
 
-func NewService(opts Options, apiOpts api.ServerOptions) (*coreService, error) {
+func NewService(cfg *setting.Cfg, opts Options, apiOpts api.ServerOptions) (*coreService, error) {
 	s := &coreService{
 		opts:    opts,
 		apiOpts: apiOpts,
+		cfg:     cfg,
 	}
 	s.BasicService = services.NewBasicService(s.start, s.running, s.stop)
 	return s, nil
 }
 
 func (s *coreService) start(_ context.Context) error {
-	serv, err := Initialize(s.cla, s.opts, s.apiOpts)
+	serv, err := Initialize(s.cfg, s.opts, s.apiOpts)
 	if err != nil {
 		return err
 	}

pkg/server/wire.go

@@ -1,6 +1,9 @@
 //go:build wireinject
 // +build wireinject
 
+// This file should contain wire sets used by both OSS and Enterprise builds.
+// Use wireext_oss.go and wireext_enterprise.go for sets that are specific to
+// the respective builds.
 package server
 
 import (
@@ -325,7 +328,6 @@ var wireBasicSet = wire.NewSet(
 	grpcserver.ProvideHealthService,
 	grpcserver.ProvideReflectionService,
 	interceptors.ProvideAuthenticator,
-	setting.NewCfgFromArgs,
 	kind.ProvideService, // The registry of known kinds
 	sqlstash.ProvideSQLEntityServer,
 	resolver.ProvideEntityReferenceResolver,
@@ -401,17 +403,31 @@ var wireTestSet = wire.NewSet(
 	wire.Bind(new(oauthtoken.OAuthTokenService), new(*oauthtokentest.Service)),
 )
 
-func Initialize(cla setting.CommandLineArgs, opts Options, apiOpts api.ServerOptions) (*Server, error) {
+func Initialize(cfg *setting.Cfg, opts Options, apiOpts api.ServerOptions) (*Server, error) {
 	wire.Build(wireExtsSet)
 	return &Server{}, nil
 }
 
-func InitializeForTest(cla setting.CommandLineArgs, opts Options, apiOpts api.ServerOptions) (*TestEnv, error) {
+func InitializeForTest(cfg *setting.Cfg, opts Options, apiOpts api.ServerOptions) (*TestEnv, error) {
 	wire.Build(wireExtsTestSet)
 	return &TestEnv{Server: &Server{}, SQLStore: &sqlstore.SQLStore{}}, nil
 }
 
-func InitializeForCLI(cla setting.CommandLineArgs) (Runner, error) {
+func InitializeForCLI(cfg *setting.Cfg) (Runner, error) {
 	wire.Build(wireExtsCLISet)
 	return Runner{}, nil
 }
+
+// InitializeForCLITarget is a simplified set of dependencies for the CLI, used
+// by the server target subcommand to launch specific dskit modules.
+func InitializeForCLITarget(cfg *setting.Cfg) (ModuleRunner, error) {
+	wire.Build(wireExtsBaseCLISet)
+	return ModuleRunner{}, nil
+}
+
+// InitializeModuleServer is a simplified set of dependencies for the CLI,
+// suitable for running background services and targeting dskit modules.
+func InitializeModuleServer(cfg *setting.Cfg, opts Options, apiOpts api.ServerOptions) (*ModuleServer, error) {
+	wire.Build(wireExtsModuleServerSet)
+	return &ModuleServer{}, nil
+}

pkg/server/wireexts_oss.go

@@ -1,10 +1,12 @@
 //go:build wireinject && oss
 // +build wireinject,oss
 
+// This file should contain wiresets which contain OSS-specific implementations.
 package server
 
 import (
 	"github.com/google/wire"
+
 	"github.com/grafana/grafana/pkg/plugins"
 	"github.com/grafana/grafana/pkg/registry"
 	"github.com/grafana/grafana/pkg/registry/backgroundsvcs"
@@ -20,6 +22,8 @@ import (
 	"github.com/grafana/grafana/pkg/services/datasources/guardian"
 	"github.com/grafana/grafana/pkg/services/encryption"
 	encryptionprovider "github.com/grafana/grafana/pkg/services/encryption/provider"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
+	"github.com/grafana/grafana/pkg/services/hooks"
 	"github.com/grafana/grafana/pkg/services/kmsproviders"
 	"github.com/grafana/grafana/pkg/services/kmsproviders/osskmsproviders"
 	"github.com/grafana/grafana/pkg/services/ldap"
@@ -103,3 +107,22 @@ var wireExtsTestSet = wire.NewSet(
 	wireTestSet,
 	wireExtsBasicSet,
 )
+
+// The wireExtsBaseCLISet is a simplified set of dependencies for the OSS CLI,
+// suitable for running background services and targeted dskit modules without
+// starting up the full Grafana server.
+var wireExtsBaseCLISet = wire.NewSet(
+	NewModuleRunner,
+
+	featuremgmt.ProvideManagerService,
+	featuremgmt.ProvideToggles,
+	hooks.ProvideService,
+	setting.ProvideProvider, wire.Bind(new(setting.Provider), new(*setting.OSSImpl)),
+	licensing.ProvideService, wire.Bind(new(licensing.Licensing), new(*licensing.OSSLicensingService)),
+)
+
+// wireModuleServerSet is a wire set for the ModuleServer.
+var wireExtsModuleServerSet = wire.NewSet(
+	NewModule,
+	wireExtsBaseCLISet,
+)

pkg/tests/testinfra/testinfra.go

@@ -45,11 +45,12 @@ func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.Tes
 	setting.IsEnterprise = extensions.IsEnterprise
 	listener, err := net.Listen("tcp", "127.0.0.1:0")
 	require.NoError(t, err)
-	cmdLineArgs := setting.CommandLineArgs{Config: cfgPath, HomePath: grafDir}
+	cfg, err := setting.NewCfgFromArgs(setting.CommandLineArgs{Config: cfgPath, HomePath: grafDir})
+	require.NoError(t, err)
 	serverOpts := server.Options{Listener: listener, HomePath: grafDir}
 	apiServerOpts := api.ServerOptions{Listener: listener}
 
-	env, err := server.InitializeForTest(cmdLineArgs, serverOpts, apiServerOpts)
+	env, err := server.InitializeForTest(cfg, serverOpts, apiServerOpts)
 	require.NoError(t, err)
 	require.NoError(t, env.SQLStore.Sync())