Skip to content

Commit

Permalink
RSDK-9440 Report machine state through GetMachineStatus (#4616)
Browse files Browse the repository at this point in the history
  • Loading branch information
@benjirewis
benjirewis authored Jan 9, 2025
1 parent 38c8eb0 commit 4775e98
Show file tree
Hide file tree
Showing 12 changed files with 585 additions and 180 deletions.
9 changes: 9 additions & 0 deletions config/config.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,12 @@ type Config struct {
// Revision contains the current revision of the config.
Revision string

// Initial represents whether this is an "initial" config passed in by web
// server entrypoint code. If true, the robot will continue to report a state
// of initializing after applying this config. If false, the robot will
// report a state of running after applying this config.
Initial bool

// toCache stores the JSON marshalled version of the config to be cached. It should be a copy of
// the config pulled from cloud with minor changes.
// This version is kept because the config is changed as it moves through the system.
Expand Down Expand Up @@ -101,6 +107,7 @@ type configData struct {
LogConfig []logging.LoggerPatternConfig `json:"log,omitempty"`
Revision string `json:"revision,omitempty"`
MaintenanceConfig *MaintenanceConfig `json:"maintenance,omitempty"`
PackagePath string `json:"package_path,omitempty"`
}

// AppValidationStatus refers to the.
Expand Down Expand Up @@ -308,6 +315,7 @@ func (c *Config) UnmarshalJSON(data []byte) error {
c.LogConfig = conf.LogConfig
c.Revision = conf.Revision
c.MaintenanceConfig = conf.MaintenanceConfig
c.PackagePath = conf.PackagePath

return nil
}
Expand Down Expand Up @@ -340,6 +348,7 @@ func (c Config) MarshalJSON() ([]byte, error) {
LogConfig: c.LogConfig,
Revision: c.Revision,
MaintenanceConfig: c.MaintenanceConfig,
PackagePath: c.PackagePath,
})
}

Expand Down
51 changes: 51 additions & 0 deletions robot/client/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,12 @@ var (

// defaultResourcesTimeout is the default timeout for getting resources.
defaultResourcesTimeout = 5 * time.Second

// DoNotWaitForRunning should be set only in tests to allow connecting to
// still-initializing machines. Note that robot clients in production (not in
// a testing environment) will already allow connecting to still-initializing
// machines.
DoNotWaitForRunning = atomic.Bool{}
)

// RobotClient satisfies the robot.Robot interface through a gRPC based
Expand Down Expand Up @@ -288,6 +294,41 @@ func New(ctx context.Context, address string, clientLogger logging.ZapCompatible
return nil, err
}

// If running in a testing environment, wait for machine to report a state of
// running. We often establish connections in tests and expect resources to
// be immediately available once the web service has started; resources will
// not be available when the machine is still initializing.
//
// It is expected that golang SDK users will handle lack of resource
// availability due to the machine being in an initializing state themselves.
//
// Allow this behavior to be turned off in some tests that specifically want
// to examine the behavior of a machine in an initializing state through the
// use of a global variable.
if testing.Testing() && !DoNotWaitForRunning.Load() {
for {
if ctx.Err() != nil {
return nil, multierr.Combine(ctx.Err(), rc.conn.Close())
}

mStatus, err := rc.MachineStatus(ctx)
if err != nil {
// Allow for MachineStatus to not be injected/implemented in some tests.
if status.Code(err) == codes.Unimplemented {
break
}
// Ignore error from Close and just return original machine status error.
utils.UncheckedError(rc.conn.Close())
return nil, err
}

if mStatus.State == robot.StateRunning {
break
}
time.Sleep(50 * time.Millisecond)
}
}

// refresh once to hydrate the robot.
if err := rc.Refresh(ctx); err != nil {
return nil, multierr.Combine(err, rc.conn.Close())
Expand Down Expand Up @@ -1115,6 +1156,16 @@ func (rc *RobotClient) MachineStatus(ctx context.Context) (robot.MachineStatus,
mStatus.Resources = append(mStatus.Resources, resStatus)
}

switch resp.State {
case pb.GetMachineStatusResponse_STATE_UNSPECIFIED:
rc.logger.CError(ctx, "received unspecified machine state")
mStatus.State = robot.StateUnknown
case pb.GetMachineStatusResponse_STATE_INITIALIZING:
mStatus.State = robot.StateInitializing
case pb.GetMachineStatusResponse_STATE_RUNNING:
mStatus.State = robot.StateRunning
}

return mStatus, nil
}

Expand Down
35 changes: 22 additions & 13 deletions robot/client/client_session_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,11 @@ func TestClientSessionOptions(t *testing.T) {
return &dummyEcho{Named: arbName.AsNamed()}, nil
},
ResourceRPCAPIsFunc: func() []resource.RPCAPI { return nil },
LoggerFunc: func() logging.Logger { return logger },
SessMgr: sessMgr,
MachineStatusFunc: func(_ context.Context) (robot.MachineStatus, error) {
return robot.MachineStatus{State: robot.StateRunning}, nil
},
LoggerFunc: func() logging.Logger { return logger },
SessMgr: sessMgr,
}

svc := web.New(injectRobot, logger)
Expand All @@ -129,13 +132,11 @@ func TestClientSessionOptions(t *testing.T) {
Disable: true,
})))
}
roboClient, err := client.New(ctx, addr, logger, opts...)
test.That(t, err, test.ShouldBeNil)

injectRobot.Mu.Lock()
injectRobot.MachineStatusFunc = func(ctx context.Context) (robot.MachineStatus, error) {
session.SafetyMonitorResourceName(ctx, someTargetName1)
return robot.MachineStatus{}, nil
return robot.MachineStatus{State: robot.StateRunning}, nil
}
injectRobot.Mu.Unlock()

Expand Down Expand Up @@ -179,6 +180,8 @@ func TestClientSessionOptions(t *testing.T) {
}
sessMgr.mu.Unlock()

roboClient, err := client.New(ctx, addr, logger, opts...)
test.That(t, err, test.ShouldBeNil)
resp, err := roboClient.MachineStatus(nextCtx)
test.That(t, err, test.ShouldBeNil)
test.That(t, resp, test.ShouldNotBeNil)
Expand Down Expand Up @@ -289,6 +292,9 @@ func TestClientSessionExpiration(t *testing.T) {
ResourceByNameFunc: func(name resource.Name) (resource.Resource, error) {
return &dummyEcho1, nil
},
MachineStatusFunc: func(_ context.Context) (robot.MachineStatus, error) {
return robot.MachineStatus{State: robot.StateRunning}, nil
},
ResourceRPCAPIsFunc: func() []resource.RPCAPI { return nil },
LoggerFunc: func() logging.Logger { return logger },
SessMgr: sessMgr,
Expand All @@ -306,8 +312,6 @@ func TestClientSessionExpiration(t *testing.T) {
Disable: true,
})))
}
roboClient, err := client.New(ctx, addr, logger, opts...)
test.That(t, err, test.ShouldBeNil)

injectRobot.Mu.Lock()
var capSessID uuid.UUID
Expand All @@ -317,7 +321,7 @@ func TestClientSessionExpiration(t *testing.T) {
panic("expected session")
}
capSessID = sess.ID()
return robot.MachineStatus{}, nil
return robot.MachineStatus{State: robot.StateRunning}, nil
}
injectRobot.Mu.Unlock()

Expand Down Expand Up @@ -372,6 +376,8 @@ func TestClientSessionExpiration(t *testing.T) {
}
sessMgr.mu.Unlock()

roboClient, err := client.New(ctx, addr, logger, opts...)
test.That(t, err, test.ShouldBeNil)
resp, err := roboClient.MachineStatus(nextCtx)
test.That(t, err, test.ShouldBeNil)
test.That(t, resp, test.ShouldNotBeNil)
Expand Down Expand Up @@ -480,8 +486,11 @@ func TestClientSessionResume(t *testing.T) {
injectRobot := &inject.Robot{
ResourceNamesFunc: func() []resource.Name { return []resource.Name{} },
ResourceRPCAPIsFunc: func() []resource.RPCAPI { return nil },
LoggerFunc: func() logging.Logger { return logger },
SessMgr: sessMgr,
MachineStatusFunc: func(_ context.Context) (robot.MachineStatus, error) {
return robot.MachineStatus{State: robot.StateRunning}, nil
},
LoggerFunc: func() logging.Logger { return logger },
SessMgr: sessMgr,
}

svc := web.New(injectRobot, logger)
Expand All @@ -496,8 +505,6 @@ func TestClientSessionResume(t *testing.T) {
Disable: true,
})))
}
roboClient, err := client.New(ctx, addr, logger, opts...)
test.That(t, err, test.ShouldBeNil)

var capMu sync.Mutex
var startCalled int
Expand Down Expand Up @@ -536,10 +543,12 @@ func TestClientSessionResume(t *testing.T) {
panic("expected session")
}
capSessID = sess.ID()
return robot.MachineStatus{}, nil
return robot.MachineStatus{State: robot.StateRunning}, nil
}
injectRobot.Mu.Unlock()

roboClient, err := client.New(ctx, addr, logger, opts...)
test.That(t, err, test.ShouldBeNil)
resp, err := roboClient.MachineStatus(nextCtx)
test.That(t, err, test.ShouldBeNil)
test.That(t, resp, test.ShouldNotBeNil)
Expand Down
Loading

0 comments on commit 4775e98

Please sign in to comment.