agent startup improvements (#865)

pkg/agent/agent.go

@@ -26,6 +26,7 @@ import (
 	"github.com/Axway/agent-sdk/pkg/config"
 	"github.com/Axway/agent-sdk/pkg/customunit"
 	"github.com/Axway/agent-sdk/pkg/util"
+	"github.com/Axway/agent-sdk/pkg/util/errors"
 	hc "github.com/Axway/agent-sdk/pkg/util/healthcheck"
 	"github.com/Axway/agent-sdk/pkg/util/log"
 )
@@ -238,7 +239,19 @@ func InitializeProfiling(cpuProfile, memProfile string) {
 	}
 }
 
+func FinalizeInitialization() error {
+	err := registerExternalIDPs()
+	if err != nil {
+		logger.WithError(err).Error("failed to register CRDs for external IdP config")
+	}
+	return nil
+}
+
 func registerExternalIDPs() error {
+	if !util.IsNotTest() || !agent.agentFeaturesCfg.ConnectionToCentralEnabled() || agent.cfg.GetUsageReportingConfig().IsOfflineMode() {
+		return nil
+	}
+
 	if agent.cfg.GetAgentType() != config.TraceabilityAgent {
 		idPCfg := agent.agentFeaturesCfg.GetExternalIDPConfig()
 		if idPCfg == nil {
@@ -261,6 +274,24 @@ func registerExternalIDPs() error {
 	return nil
 }
 
+func CacheInitSync() error {
+	if !util.IsNotTest() || !agent.agentFeaturesCfg.ConnectionToCentralEnabled() || agent.cfg.GetUsageReportingConfig().IsOfflineMode() {
+		return nil
+	}
+
+	eventSync, err := newEventSync()
+	if err != nil {
+		return errors.Wrap(errors.ErrInitServicesNotReady, err.Error())
+	}
+
+	if err := eventSync.SyncCache(); err != nil {
+		return errors.Wrap(errors.ErrInitServicesNotReady, err.Error())
+	}
+	// set the rebuild function in the agent resource manager
+	agent.agentResourceManager.SetRebuildCacheFunc(eventSync)
+	return nil
+}
+
 func registerCredentialProvider(idp config.IDPConfig, tlsCfg config.TLSConfig, proxyURL string, clientTimeout time.Duration) error {
 	err := GetAuthProviderRegistry().RegisterProvider(idp, tlsCfg, proxyURL, clientTimeout)
 	if err != nil {

pkg/agent/eventsync.go

@@ -24,8 +24,8 @@ type EventSync struct {
 	discoveryCache *discoveryCache
 }
 
-// NewEventSync creates an EventSync
-func NewEventSync() (*EventSync, error) {
+// newEventSync creates an EventSync
+func newEventSync() (*EventSync, error) {
 	migrations := []migrate.Migrator{}
 
 	// Make sure only DA and Governance agents run migration processes
@@ -84,12 +84,7 @@ func (es *EventSync) SyncCache() error {
 		}
 	}
 
-	err := registerExternalIDPs()
-	if err != nil {
-		logger.WithError(err).Warn("failed to register CRDs for external IdP config")
-	}
-
-	err = es.startCentralEventProcessor()
+	err := es.startCentralEventProcessor()
 	if err != nil {
 		return err
 	}

pkg/agent/eventsync_test.go

@@ -44,7 +44,7 @@ func TestEventSync_pollMode(t *testing.T) {
 	InitializeForTest(mc)
 	assert.Nil(t, err)
 
-	es, err := NewEventSync()
+	es, err := newEventSync()
 	assert.Nil(t, err)
 	assert.NotNil(t, es.watchTopic)
 	assert.NotNil(t, es.discoveryCache)
@@ -95,7 +95,7 @@ func TestEventSync_streamMode(t *testing.T) {
 	InitializeForTest(mc)
 	assert.Nil(t, err)
 
-	es, err := NewEventSync()
+	es, err := newEventSync()
 	assert.Nil(t, err)
 	assert.NotNil(t, es.watchTopic)
 	assert.NotNil(t, es.discoveryCache)

pkg/cmd/root.go

@@ -37,11 +37,14 @@ const (
 	httpprofile            = "httpprofile"
 )
 
+type NewCommandOption func(*agentRootCommand)
+
 // CommandHandler - Root command execution handler
 type CommandHandler func() error
 
 // InitConfigHandler - Handler to be invoked on config initialization
 type InitConfigHandler func(centralConfig config.CentralConfig) (interface{}, error)
+type FinalizeAgentInitHandler func() error
 
 // AgentRootCmd - Root Command for the Agents
 type AgentRootCmd interface {
@@ -61,6 +64,7 @@ type agentRootCommand struct {
 	rootCmd           *cobra.Command
 	commandHandler    CommandHandler
 	initConfigHandler InitConfigHandler
+	finalizeAgentInit FinalizeAgentInitHandler
 	agentType         config.AgentType
 	props             properties.Properties
 	statusCfg         config.StatusConfig
@@ -95,8 +99,14 @@ func buildAgentInfo(desc string) string {
 	return fmt.Sprintf("%s version %s-%s, Amplify Agents SDK version %s", desc, BuildVersion, BuildCommitSha, SDKBuildVersion)
 }
 
+func WithFinalizeAgentInitFunc(f FinalizeAgentInitHandler) NewCommandOption {
+	return func(c *agentRootCommand) {
+		c.finalizeAgentInit = f
+	}
+}
+
 // NewRootCmd - Creates a new Agent Root Command
-func NewRootCmd(exeName, desc string, initConfigHandler InitConfigHandler, commandHandler CommandHandler, agentType config.AgentType) AgentRootCmd {
+func NewRootCmd(exeName, desc string, initConfigHandler InitConfigHandler, commandHandler CommandHandler, agentType config.AgentType, opts ...NewCommandOption) AgentRootCmd {
 	c := &agentRootCommand{
 		agentName:         exeName,
 		commandHandler:    commandHandler,
@@ -106,6 +116,10 @@ func NewRootCmd(exeName, desc string, initConfigHandler InitConfigHandler, comma
 		initialized:       false,
 	}
 
+	for _, o := range opts {
+		o(c)
+	}
+
 	// use the description from the build if available
 	if BuildAgentDescription != "" {
 		desc = BuildAgentDescription
@@ -136,7 +150,7 @@ func NewRootCmd(exeName, desc string, initConfigHandler InitConfigHandler, comma
 }
 
 // NewCmd - Creates a new Agent Root Command using existing cmd
-func NewCmd(rootCmd *cobra.Command, exeName, desc string, initConfigHandler InitConfigHandler, commandHandler CommandHandler, agentType config.AgentType) AgentRootCmd {
+func NewCmd(rootCmd *cobra.Command, exeName, desc string, initConfigHandler InitConfigHandler, commandHandler CommandHandler, agentType config.AgentType, opts ...NewCommandOption) AgentRootCmd {
 	c := &agentRootCommand{
 		agentName:         exeName,
 		commandHandler:    commandHandler,
@@ -146,6 +160,10 @@ func NewCmd(rootCmd *cobra.Command, exeName, desc string, initConfigHandler Init
 		initialized:       false,
 	}
 
+	for _, o := range opts {
+		o(c)
+	}
+
 	// use the description from the build if available
 	if BuildAgentDescription != "" {
 		desc = BuildAgentDescription
@@ -354,18 +372,21 @@ func (c *agentRootCommand) initConfig() error {
 }
 
 func (c *agentRootCommand) finishInit() error {
-	if util.IsNotTest() && c.agentFeaturesCfg.ConnectionToCentralEnabled() && !c.centralCfg.GetUsageReportingConfig().IsOfflineMode() {
-		eventSync, err := agent.NewEventSync()
-		if err != nil {
-			return errors.Wrap(errors.ErrInitServicesNotReady, err.Error())
-		}
+	err := agent.FinalizeInitialization()
+	if err != nil {
+		return err
+	}
 
-		if err := eventSync.SyncCache(); err != nil {
-			return errors.Wrap(errors.ErrInitServicesNotReady, err.Error())
+	if c.finalizeAgentInit != nil {
+		err := c.finalizeAgentInit()
+		if err != nil {
+			return err
 		}
-		// set the rebuild function in the agent resource manager
-		agent.GetAgentResourceManager().SetRebuildCacheFunc(eventSync)
+	}
 
+	err = agent.CacheInitSync()
+	if err != nil {
+		return err
 	}
 
 	// Start the initial and recurring version check jobs