Skip to content

Commit

Permalink
Refactor command-running to run oneshots after startup
Browse files Browse the repository at this point in the history
This refactors the command-running utilities to take the fx.Options as
variadic arguments, so that fxapps.OneShot can add its own fx.Option.

It uses a bit of reflection to convince Fx to gather the arguments to a
function without actually calling that function immediately, allowing
the oneShot function to run after `app.Start` has finished, when we're
sure that all components have started.  This works around the startup
ordering bug in uber-go/fx#918.
  • Loading branch information
djmitche committed Aug 12, 2022
1 parent 0cf5589 commit f031214
Show file tree
Hide file tree
Showing 8 changed files with 97 additions and 47 deletions.
6 changes: 2 additions & 4 deletions cmd/agent/flare/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@ import (
"github.com/DataDog/dd-agent-comp-experiments/cmd/common"
"github.com/DataDog/dd-agent-comp-experiments/comp/core/flare"
"github.com/DataDog/dd-agent-comp-experiments/comp/core/ipc/ipcclient"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/fxapps"
"github.com/spf13/cobra"
"go.uber.org/fx"
)

var (
Expand All @@ -27,11 +27,9 @@ var (
)

func command(_ *cobra.Command, args []string) error {
app := fx.New(
return fxapps.OneShot(flareCmd,
common.SharedOptions(root.ConfFilePath, true),
common.OneShot(flareCmd),
)
return common.RunApp(app)
}

func getFlareRemote(ipcclient ipcclient.Component) (string, error) {
Expand Down
6 changes: 2 additions & 4 deletions cmd/agent/health/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,8 @@ import (
"github.com/DataDog/dd-agent-comp-experiments/cmd/common"
"github.com/DataDog/dd-agent-comp-experiments/comp/core/health"
"github.com/DataDog/dd-agent-comp-experiments/comp/core/ipc/ipcclient"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/fxapps"
"github.com/spf13/cobra"
"go.uber.org/fx"
)

var (
Expand All @@ -26,11 +26,9 @@ var (
)

func command(_ *cobra.Command, args []string) error {
app := fx.New(
return fxapps.OneShot(healthCmd,
common.SharedOptions(root.ConfFilePath, true),
common.OneShot(healthCmd),
)
return common.RunApp(app)
}

func getHealthRemote(ipcclient ipcclient.Component) (map[string]health.ComponentHealth, error) {
Expand Down
4 changes: 2 additions & 2 deletions cmd/agent/run/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import (
"github.com/DataDog/dd-agent-comp-experiments/comp/logs"
"github.com/DataDog/dd-agent-comp-experiments/comp/logs/launchers/file"
"github.com/DataDog/dd-agent-comp-experiments/comp/trace"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/fxapps"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/startup"
"github.com/spf13/cobra"
"go.uber.org/fx"
Expand All @@ -35,7 +36,7 @@ func logsAgentPluginOptions() fx.Option {
}

func run(_ *cobra.Command, args []string) error {
app := fx.New(
return fxapps.Run(
common.SharedOptions(root.ConfFilePath, false),
fx.Supply(logs.BundleParams{
AutoStart: startup.IfConfigured,
Expand All @@ -47,5 +48,4 @@ func run(_ *cobra.Command, args []string) error {
trace.Bundle,
logsAgentPluginOptions(),
)
return common.RunApp(app)
}
20 changes: 11 additions & 9 deletions cmd/agent/status/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"github.com/DataDog/dd-agent-comp-experiments/cmd/common"
"github.com/DataDog/dd-agent-comp-experiments/comp/core/ipc/ipcclient"
"github.com/DataDog/dd-agent-comp-experiments/comp/core/status"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/fxapps"
"github.com/spf13/cobra"
"go.uber.org/fx"
)
Expand All @@ -26,19 +27,20 @@ var (
}
)

type cmdArgs struct {
section string
}

func command(_ *cobra.Command, args []string) error {
var section string
var cmdArgs cmdArgs
if len(args) > 0 {
section = args[0]
cmdArgs.section = args[0]
}

app := fx.New(
return fxapps.OneShot(statusCmd,
fx.Supply(cmdArgs),
common.SharedOptions(root.ConfFilePath, true),
common.OneShot(func(ipcclient ipcclient.Component, status status.Component) error {
return statusCmd(ipcclient, status, section)
}),
)
return common.RunApp(app)
}

func getStatusRemote(ipcclient ipcclient.Component, section string) (string, error) {
Expand All @@ -56,8 +58,8 @@ func getStatusRemote(ipcclient ipcclient.Component, section string) (string, err
return content["status"], nil
}

func statusCmd(ipcclient ipcclient.Component, status status.Component, section string) error {
statusStr, err := getStatusRemote(ipcclient, section)
func statusCmd(ipcclient ipcclient.Component, status status.Component, cmdArgs cmdArgs) error {
statusStr, err := getStatusRemote(ipcclient, cmdArgs.section)
if err != nil {
return err
}
Expand Down
22 changes: 0 additions & 22 deletions cmd/common/oneshot.go

This file was deleted.

4 changes: 2 additions & 2 deletions cmd/traceagent/run/command.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"github.com/DataDog/dd-agent-comp-experiments/cmd/agent/root"
"github.com/DataDog/dd-agent-comp-experiments/cmd/common"
"github.com/DataDog/dd-agent-comp-experiments/comp/trace"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/fxapps"
"github.com/DataDog/dd-agent-comp-experiments/pkg/util/startup"
"github.com/spf13/cobra"
"go.uber.org/fx"
Expand All @@ -25,12 +26,11 @@ var (
)

func run(_ *cobra.Command, args []string) error {
app := fx.New(
return fxapps.Run(
common.SharedOptions(root.ConfFilePath, false),
fx.Supply(trace.BundleParams{
AutoStart: startup.IfConfigured,
}),
trace.Bundle,
)
return common.RunApp(app)
}
73 changes: 73 additions & 0 deletions pkg/util/fxapps/oneshot.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.

package fxapps

import (
"context"
"reflect"

"go.uber.org/fx"
)

var errorInterface = reflect.TypeOf((*error)(nil)).Elem()

// OneShot runs the given function in an fx.App using the supplied options.
// The function's arguments are supplied by Fx and can be any provided type.
// The function must return `error.
//
// The resulting app starts all components, then invokes the function, then
// immediately shuts down. This is typically used for command-line tools like
// `agent status`.
func OneShot(oneShotFunc interface{}, opts ...fx.Option) error {
ftype := reflect.TypeOf(oneShotFunc)
if ftype == nil || ftype.Kind() != reflect.Func {
panic("OneShot requires a function as its first argument")
}

// verify it returns error
if ftype.NumOut() != 1 || !ftype.Out(0).Implements(errorInterface) {
panic("OneShot function must return error or nothing")
}

// build an function with the same signature as oneShotFunc that will
// capture the args and do nothing.
var oneShotArgs []reflect.Value
captureArgs := reflect.MakeFunc(
ftype,
func(args []reflect.Value) []reflect.Value {
oneShotArgs = args
// return a single nil value of type error
return []reflect.Value{reflect.Zero(errorInterface)}
})
// fx.Invoke that function to capture the args at startup
opt := fx.Invoke(captureArgs.Interface())
opts = append(opts, opt)
app := fx.New(opts...)

startCtx, cancel := context.WithTimeout(context.Background(), app.StartTimeout())
defer cancel()

if err := app.Start(startCtx); err != nil {
return err
}

// call the original oneShotFunc with the args captured during
// app startup
res := reflect.ValueOf(oneShotFunc).Call(oneShotArgs)
if !res[0].IsNil() {
err := res[0].Interface().(error)
return err
}

stopCtx, cancel := context.WithTimeout(context.Background(), app.StopTimeout())
defer cancel()

if err := app.Stop(stopCtx); err != nil {
return err
}

return nil
}
9 changes: 5 additions & 4 deletions cmd/common/run.go → pkg/util/fxapps/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,18 @@
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2016-present Datadog, Inc.

package common
package fxapps

import (
"context"

"go.uber.org/fx"
)

// RunApp is similar to fx.App#Run, but returns an error or nil when the app
// completes, instead of exiting the process.
func RunApp(app *fx.App) error {
// Run runs an fx.App using the supplied options, returning any errors.
func Run(opts ...fx.Option) error {
app := fx.New(opts...)

startCtx, cancel := context.WithTimeout(context.Background(), app.StartTimeout())
defer cancel()

Expand Down

0 comments on commit f031214

Please sign in to comment.