Add code to handle disable actions and support for merging HGAP immediate GS with local GS cache (#53)

* Create go.yml

* initial changes to support disabling a command

* Print payload

* update bits

* update bits

* Only report not empty statuses

* Only report not empty statuses

* Only report not empty statuses

* Update bits

* Update bits

* Update bits

* Update bits

* Update bits

* Update status to uploadf

* Update status to uploadf

* update bits

* Add extension name in status

* add logging for reporting status

* change the format of the status to upload

* change the format of the status to upload

* change the format of the status to upload

* update bits

* Fix unit tests

* Add notification when the extension is deleted

* Add notification when the extension is deleted

* Handle non-modified vmsettings

* Remove unused bits

* Fix comment

* Fix unit test

* Add fake server for other unit tests

* Reset seqnum when disabling extensions

* Add code to store in the cache goal states in terminal states

* Remove status after it's beeing store

* Add some logging

* Add logs

* Remove disable goal states after they were processed

* Report goal states in terminal status to HGAP

* Remove disable status from local cache

* Add logging

* Change the check for when the extension is disabled

* Fix issue of not reporting back updated goal states

* Change the frequency to report the status of a command from 30 to 15 seconds

* remove unused error

* Change file extension of immediate goal states to make sure it gets copied when upgrading the extension

* Test logs

* Add changes to immediately storing final states in local file

* Use correct StatusType strings

* Lock status file to make sure it keeps consistent

* Fix unit tests

Author: fifuente

internal/cmds/cmds.go

@@ -20,6 +20,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/appendblob"
 	"github.com/Azure/azure-sdk-for-go/storage"
 	"github.com/Azure/run-command-handler-linux/internal/cleanup"
+	"github.com/Azure/run-command-handler-linux/internal/commandProcessor"
 	"github.com/Azure/run-command-handler-linux/internal/constants"
 	"github.com/Azure/run-command-handler-linux/internal/exec"
 	"github.com/Azure/run-command-handler-linux/internal/files"
@@ -39,7 +40,7 @@ import (
 
 const (
 	maxScriptSize         = 256 * 1024
-	updateStatusInSeconds = 30
+	updateStatusInSeconds = 15
 )
 
 const (
@@ -86,13 +87,17 @@ func update(ctx *log.Context, h types.HandlerEnvironment, report *types.RunComma
 }
 
 func disable(ctx *log.Context, h types.HandlerEnvironment, report *types.RunCommandInstanceView, metadata types.RCMetadata, c types.Cmd) (string, string, error, int) {
-	exitCode, err := immediatecmds.Disable(ctx, h, metadata.ExtName, metadata.SeqNum)
-	if err != nil {
-		return "", "", err, exitCode
+	extensionHandlerName := commandProcessor.GetExtensionName(ctx)
+	if extensionHandlerName == constants.ImmediateRunCommandHandlerName {
+		exitCode, err := immediatecmds.Disable(ctx, h, metadata.ExtName, metadata.SeqNum)
+		if err != nil {
+			return "", "", err, exitCode
+		}
 	}
 
 	ctx.Log("event", "disable")
 	pid.KillPreviousExtension(ctx, metadata.PidFilePath)
+	resetSeqNum(ctx, metadata.MostRecentSequence)
 	return "", "", nil, constants.ExitCode_Okay
 }
 
@@ -335,10 +340,16 @@ func checkAndSaveSeqNum(ctx log.Logger, seq int, mrseqPath string) (shouldExit b
 	return false, nil
 }
 
+// resetSeqNum deletes the seqNum file to reset the sequence number
+func resetSeqNum(ctx log.Logger, mrseqPath string) {
+	ctx.Log("event", "resetting seqnum by deleting file", "path", mrseqPath)
+	os.Remove(mrseqPath)
+}
+
 // Copy state of the extension from old version to new version during update (.mrseq files, .status files)
 func CopyStateForUpdate(ctx log.Logger) error {
 	// Copy .mrseq files (Most Recently executed Sequence number) that helps determine whether a sequence number of Run Command has been previously executed or not.
-	mrseqFilesNameList, mrseqFileCopyErr := copyFiles(ctx, ".mrseq", "")
+	mrseqFilesNameList, mrseqFileCopyErr := copyFiles(ctx, constants.MrSeqFileExtension, "")
 	if mrseqFileCopyErr != nil {
 		return mrseqFileCopyErr
 	}
@@ -536,7 +547,7 @@ func createDummyStatusFilesIfNeeded(ctx log.Logger, mrseqFilesNameList *list.Lis
 				continue
 			}
 
-			statusReport := types.NewStatusReport(types.StatusSuccess, "Enable", instanceViewMessage)
+			statusReport := types.NewStatusReport(types.StatusSuccess, "Enable", instanceViewMessage, extensionName)
 			rootStatusJson, err = status.MarshalStatusReportIntoJson(statusReport, true)
 			if err != nil {
 				errorMessage = fmt.Sprintf("failed to marshal status report into json for status file '%s' with error '%s'", statusFilePath, err.Error())

internal/commandProcessor/commandProcessor.go

@@ -121,7 +121,7 @@ func getRequiredInitialVariables(ctx *log.Context) (types.HandlerEnvironment, st
 		return hEnv, extensionName, seqNum, errors.Wrap(err, "failed to parse handlerEnv")
 	}
 
-	extensionName = getExtensionName(ctx)
+	extensionName = GetExtensionName(ctx)
 	seqNum, err = getSeqNum(&ctx, hEnv, extensionName)
 	if err != nil {
 		return hEnv, extensionName, seqNum, errors.Wrap(err, "failed to get seqNum")
@@ -162,7 +162,7 @@ func getHandlerEnv(ctx *log.Context) (types.HandlerEnvironment, error) {
 	return hEnv, nil
 }
 
-func getExtensionName(ctx *log.Context) string {
+func GetExtensionName(ctx *log.Context) string {
 	extensionName := os.Getenv(constants.ConfigExtensionNameEnvName)
 	ctx.Log("extensionName", extensionName)
 	return extensionName

internal/commandProcessor/commandProcessor_test.go

@@ -112,7 +112,7 @@ func Test_GetExtNameFromEnvVariable(t *testing.T) {
 	extName := "testExtension"
 	os.Setenv(constants.ConfigExtensionNameEnvName, extName)
 
-	actualExtName := getExtensionName(ctx)
+	actualExtName := GetExtensionName(ctx)
 	require.Equal(t, extName, actualExtName)
 }
 

internal/constants/constants.go

@@ -25,6 +25,9 @@ const (
 
 	StatusFileDirectory = "status"
 
+	// The directory where the immediate run command status that have reached the terminal status are stored.
+	ImmediateStatusFileDirectory = "status"
+
 	// General failed exit code when extension provisioning fails due to service errors.
 	FailedExitCodeGeneral = -1
 
@@ -51,4 +54,13 @@ const (
 	// The path of the extension in the VM with full name. This value is provided by the agent for all commands.
 	// See more in: https://github.com/Azure/azure-vmextension-publishing/wiki/2.0-Partner-Guide-Handler-Design-Details#236-summary
 	ExtensionPathEnvName = "AZURE_GUEST_AGENT_EXTENSION_PATH"
+
+	// The name of the immediate run command service
+	ImmediateRunCommandHandlerName = "runCommandService"
+
+	// The time to wait between each poll of the goal states
+	PolingIntervalInSeconds = 1
+
+	// The name of the file that contains the immediate goal states that reached the terminal status
+	ImmediateGoalStatesInTerminalStatusFileName = "immediateGoalStatesInTerminalStatusFile.status"
 )

internal/goalstate/goalstate.go

@@ -19,9 +19,15 @@ import (
 
 const (
 	enableCommand             string = "enable"
+	disableCommand            string = "disable"
 	maxExecutionTimeInMinutes int32  = 90
 )
 
+var statusToCommandMap = map[string]string{
+	"enabled":  enableCommand,
+	"disabled": disableCommand,
+}
+
 // HandleImmediateGoalState handles the immediate goal state by executing the command and waiting for it to finish.
 // ctx: The logger context.
 // setting: The settings for the command.
@@ -32,12 +38,14 @@ func HandleImmediateGoalState(ctx *log.Context, setting settings.SettingsCommon,
 	err := make(chan error)
 	go startAsync(ctx, setting, notifier, done, err)
 	select {
-	case <-err:
-		return constants.ExitCode_ImmediateTaskFailed, errors.Wrapf(<-err, "error when trying to execute goal state")
+	case e := <-err:
+		ctx.Log("error", fmt.Sprintf("error when trying to execute goal state: %v", e))
+		return constants.ExitCode_ImmediateTaskFailed, errors.Wrapf(e, "error when trying to execute goal state")
 	case <-done:
 		ctx.Log("message", "goal state successfully finished")
 		return constants.ExitCode_Okay, nil
 	case <-time.After(time.Minute * time.Duration(maxExecutionTimeInMinutes)):
+		ctx.Log("message", "timeout when trying to execute goal state")
 		return constants.ExitCode_ImmediateTaskTimeout, errors.New("timeout when trying to execute goal state")
 	}
 }
@@ -52,9 +60,11 @@ func ReportFinalStatusForImmediateGoalState(ctx *log.Context, notifier *observer
 		return errors.New("notifier is nil. Cannot report status to HGAP")
 	}
 
-	cmd, ok := commands.Cmds[enableCommand]
+	extensionState := goalStateKey.RuntimeSettingsState
+	cmdToReport := statusToCommandMap[extensionState]
+	cmd, ok := commands.Cmds[cmdToReport]
 	if !ok {
-		return errors.New("missing enable command")
+		return errors.New(fmt.Sprintf("missing command %v", extensionState))
 	}
 
 	if !cmd.ShouldReportStatus {
@@ -67,7 +77,7 @@ func ReportFinalStatusForImmediateGoalState(ctx *log.Context, notifier *observer
 		return errors.Wrapf(err, "failed to marshal instance view")
 	}
 
-	statusItem, err := status.GetSingleStatusItem(ctx, statusType, cmd, string(msg))
+	statusItem, err := status.GetSingleStatusItem(ctx, statusType, cmd, string(msg), goalStateKey.ExtensionName)
 	if err != nil {
 		return errors.Wrap(err, "failed to get status item")
 	}
@@ -88,29 +98,34 @@ func startAsync(ctx *log.Context, setting settings.SettingsCommon, notifier *obs
 		return
 	}
 
-	cmd, ok := commands.Cmds[enableCommand]
+	extensionState := *setting.ExtensionState
+	ctx.Log("message", fmt.Sprintf("starting command for extension state %v", extensionState))
+
+	cmdToReport := statusToCommandMap[extensionState]
+	cmd, ok := commands.Cmds[cmdToReport]
 	if !ok {
-		err <- errors.New("missing enable command")
+		err <- errors.New(fmt.Sprintf("missing command %v", extensionState))
 		return
 	}
 
 	// Overwrite function to report status to HGAP. This function prepares the status to be sent to the HGAP and then calls the notifier to send it.
 	cmd.Functions.ReportStatus = func(ctx *log.Context, _ types.HandlerEnvironment, metadata types.RCMetadata, statusType types.StatusType, c types.Cmd, msg string) error {
 		if !c.ShouldReportStatus {
-			ctx.Log("status", "not reported for operation (by design)")
+			ctx.Log("status", fmt.Sprintf("status not reported for operation %v (by design)", c.Name))
 			return nil
 		}
 
-		statusItem, err := status.GetSingleStatusItem(ctx, statusType, c, msg)
+		statusItem, err := status.GetSingleStatusItem(ctx, statusType, c, msg, metadata.ExtName)
 		if err != nil {
 			return errors.Wrap(err, "failed to get status item")
 		}
 
 		ctx.Log("message", fmt.Sprintf("reporting status by notifying the observer to then send to HGAP for extension name %v and seq number %v", metadata.ExtName, metadata.SeqNum))
 		return notifier.Notify(types.StatusEventArgs{
 			StatusKey: types.GoalStateKey{
-				ExtensionName: metadata.ExtName,
-				SeqNumber:     metadata.SeqNum,
+				ExtensionName:        metadata.ExtName,
+				SeqNumber:            metadata.SeqNum,
+				RuntimeSettingsState: extensionState,
 			},
 			TopLevelStatus: statusItem,
 		})

internal/goalstate/goalstate_test.go

@@ -18,6 +18,16 @@ import (
 func Test_handleSkippedImmediateGoalState_NotifyObserver(t *testing.T) {
 	ctx := log.NewContext(log.NewSyncLogger(log.NewLogfmtLogger(os.Stdout))).With("time", log.DefaultTimestamp)
 
+	// Create a temporary directory for the test
+	tempDir, err := os.MkdirTemp("", "")
+	defer os.RemoveAll(tempDir) // Clean up the temp directory after the test
+	require.NoError(t, err, "Failed to create temp directory")
+
+	err = os.Mkdir(tempDir+"/status", 0755)
+	require.NoError(t, err, "Failed to create status directory")
+
+	os.Setenv(constants.ExtensionPathEnvName, tempDir)
+
 	ctx.Log("msg", "Creating status observer")
 	obs := &status.StatusObserver{}
 	obs.Initialize(ctx)
@@ -28,8 +38,9 @@ func Test_handleSkippedImmediateGoalState_NotifyObserver(t *testing.T) {
 	notifier.Register(obs)
 
 	goalStateKey := types.GoalStateKey{
-		ExtensionName: "test",
-		SeqNumber:     1,
+		ExtensionName:        "test",
+		SeqNumber:            1,
+		RuntimeSettingsState: "enabled",
 	}
 
 	errorMsg := "Test error message"
@@ -43,15 +54,22 @@ func Test_handleSkippedImmediateGoalState_NotifyObserver(t *testing.T) {
 		EndTime:          time.Now().UTC().Format(time.RFC3339),
 	}
 
-	err := ReportFinalStatusForImmediateGoalState(ctx, notifier, goalStateKey, types.StatusSkipped, &instView)
+	err = ReportFinalStatusForImmediateGoalState(ctx, notifier, goalStateKey, types.StatusSkipped, &instView)
 	require.Nil(t, err, "HandleSkippedImmediateGoalState should not return an error")
 
 	ctx.Log("msg", "Unregistering observer")
 	notifier.Unregister()
 
 	ctx.Log("msg", "Check that the status item was received by the observer")
-	latestStatus, ok := obs.GetStatusForKey(goalStateKey)
-	require.True(t, ok, "Status item should be found")
+	_, ok := obs.GetStatusForKey(goalStateKey)
+	require.False(t, ok, "Status item should not be found because it is in a terminal state and saved locally")
+
+	// Get the status item from the local file
+	items, err := status.GetGoalStatesInTerminalStatus(ctx)
+	require.Nil(t, err, "GetGoalStatesInTerminalStatus should not return an error")
+	require.Len(t, items, 1, "There should be one status item in the local file")
+
+	latestStatus := items[0]
 	require.Equal(t, "Enable", latestStatus.Status.Operation, "Operation should be equal")
 	require.Equal(t, types.StatusSkipped, latestStatus.Status.Status, "Status should be equal")
 

internal/immediateruncommand/immediateruncommand.go

@@ -42,6 +42,7 @@ func StartImmediateRunCommand(ctx *log.Context) error {
 	communicator := hostgacommunicator.NewHostGACommunicator(vmRequestManager)
 	goalStateEventObserver.Initialize(ctx)
 
+	ctx.Log("message", fmt.Sprintf("Polling for goal state every %v seconds", constants.PolingIntervalInSeconds))
 	for {
 		ctx.Log("message", "processing new immediate run command goal states. Last processed ETag: "+lastProcessedETag)
 		newProcessedETag, err := processImmediateRunCommandGoalStates(ctx, communicator, lastProcessedETag)
@@ -52,6 +53,7 @@ func StartImmediateRunCommand(ctx *log.Context) error {
 			time.Sleep(time.Second * time.Duration(5))
 		} else {
 			lastProcessedETag = newProcessedETag
+			time.Sleep(time.Second * time.Duration(constants.PolingIntervalInSeconds))
 		}
 	}
 }
@@ -69,10 +71,17 @@ func processImmediateRunCommandGoalStates(ctx *log.Context, communicator hostgac
 		return newEtag, errors.Wrapf(err, "could not retrieve goal states for immediate run command")
 	}
 
+	// VM Settings have not changed and we should not process any new goal states
+	if newEtag == lastProcessedETag {
+		return newEtag, nil
+	}
+
 	var goalStateKeys []types.GoalStateKey
 	for _, s := range goalStates {
 		for _, setting := range s.Settings {
-			goalStateKeys = append(goalStateKeys, types.GoalStateKey{ExtensionName: *setting.ExtensionName, SeqNumber: *setting.SeqNo})
+			if setting.ExtensionState != nil {
+				goalStateKeys = append(goalStateKeys, types.GoalStateKey{ExtensionName: *setting.ExtensionName, SeqNumber: *setting.SeqNo, RuntimeSettingsState: *setting.ExtensionState})
+			}
 		}
 	}
 	goalStateEventObserver.RemoveProcessedGoalStates(goalStateKeys)
@@ -90,7 +99,7 @@ func processImmediateRunCommandGoalStates(ctx *log.Context, communicator hostgac
 				executingTasks.Increment()
 
 				ctx.Log("message", "adding goal state to the event map")
-				statusKey := types.GoalStateKey{ExtensionName: *state.ExtensionName, SeqNumber: *state.SeqNo}
+				statusKey := types.GoalStateKey{ExtensionName: *state.ExtensionName, SeqNumber: *state.SeqNo, RuntimeSettingsState: *state.ExtensionState}
 				defaultTopStatus := types.StatusItem{}
 				status := types.StatusEventArgs{TopLevelStatus: defaultTopStatus, StatusKey: statusKey}
 
@@ -116,7 +125,7 @@ func processImmediateRunCommandGoalStates(ctx *log.Context, communicator hostgac
 						StartTime:        startTime,
 						EndTime:          time.Now().UTC().Format(time.RFC3339),
 					}
-					goalstate.ReportFinalStatusForImmediateGoalState(ctx, notifier, statusKey, types.StatusSkipped, &instView)
+					goalstate.ReportFinalStatusForImmediateGoalState(ctx, notifier, statusKey, types.StatusError, &instView)
 
 				}
 			}(newGoalStates[idx])
@@ -130,7 +139,7 @@ func processImmediateRunCommandGoalStates(ctx *log.Context, communicator hostgac
 	if len(skippedGoalStates) > 0 {
 		ctx.Log("message", fmt.Sprintf("skipped %v goal states due to reaching the maximum concurrent tasks", len(skippedGoalStates)))
 		for _, skippedGoalState := range skippedGoalStates {
-			statusKey := types.GoalStateKey{ExtensionName: *skippedGoalState.ExtensionName, SeqNumber: *skippedGoalState.SeqNo}
+			statusKey := types.GoalStateKey{ExtensionName: *skippedGoalState.ExtensionName, SeqNumber: *skippedGoalState.SeqNo, RuntimeSettingsState: *skippedGoalState.ExtensionState}
 			notifier := &observer.Notifier{}
 			notifier.Register(&goalStateEventObserver)
 
@@ -165,7 +174,7 @@ func getGoalStatesToProcess(goalStates []hostgacommunicator.ImmediateExtensionGo
 
 		if validSignature {
 			for _, s := range el.Settings {
-				statusKey := types.GoalStateKey{ExtensionName: *s.ExtensionName, SeqNumber: *s.SeqNo}
+				statusKey := types.GoalStateKey{ExtensionName: *s.ExtensionName, SeqNumber: *s.SeqNo, RuntimeSettingsState: *s.ExtensionState}
 				_, goalStateAlreadyProcessed := goalStateEventObserver.GetStatusForKey(statusKey)
 				if !goalStateAlreadyProcessed {
 					if len(newGoalStates) < maxTasksToFetch {

internal/observer/observer.go

@@ -7,4 +7,6 @@ import "github.com/Azure/run-command-handler-linux/internal/types"
 type Observer interface {
 	// OnNotify is called when the status changes
 	OnNotify(types.StatusEventArgs) error
+	// OnDemandNotify is called when the observer needs to report the status immediately
+	OnDemandNotify() error
 }