Implementing New Sequence Number Management and Fixing how we get the extension Sequence Number

.gitignore

@@ -28,3 +28,4 @@ _testmain.go
 *.exe
 *.test
 *.prof
+mrseq

go.mod

@@ -12,7 +12,10 @@ require (
 	github.com/xeipuuv/gojsonschema v0.0.0-20160623135812-c539bca196be
 )
 
-require github.com/go-kit/log v0.2.0
+require (
+	github.com/go-kit/log v0.2.0
+	go.uber.org/mock v0.4.0
+)
 
 require (
 	github.com/cilium/ebpf v0.9.1 // indirect

go.sum

@@ -51,6 +51,8 @@ github.com/xeipuuv/gojsonschema v0.0.0-20160623135812-c539bca196be h1:sRGd3e18iz
 github.com/xeipuuv/gojsonschema v0.0.0-20160623135812-c539bca196be/go.mod h1:5yf86TLmAcydyeJq5YvxkGPE2fm/u4myDekKRoLuqhs=
 go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
 go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
+go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
+go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=
 golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.2.0 h1:ljd4t30dBnAvMZaQCevtY0xLLD0A+bRZXbgLMLU1F/A=
 golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=

internal/seqno/seqno.go

@@ -0,0 +1,64 @@
+package seqno
+
+import (
+	"github.com/Azure/applicationhealth-extension-linux/pkg/logging"
+	"github.com/Azure/azure-extension-platform/pkg/extensionerrors"
+	"github.com/Azure/azure-extension-platform/pkg/seqno"
+	"github.com/go-kit/log"
+)
+
+type SequenceNumberManager interface {
+	// GetCurrentSequenceNumber returns the current sequence number the extension is using
+	GetCurrentSequenceNumber(el log.Logger, name, version string) (uint, error)
+
+	// GetSequenceNumber retrieves the sequence number from the MRSEQ file
+	GetSequenceNumber(name, version string) (uint, error)
+
+	// SetSequenceNumber sets the sequence number to the MRSEQ file.
+	SetSequenceNumber(name, version string, seqNo uint) error
+
+	// FindSeqNum returns the requested the sequence number from either the environment variable or
+	// the most recently used file under the config folder.
+	// Note that this is different than just choosing the highest number, which may be incorrect
+	FindSeqNum(configFolder string) (uint, error)
+}
+
+type SeqNumManager struct {
+}
+
+func (s *SeqNumManager) GetSequenceNumber(name string, version string) (uint, error) {
+	retriever := &seqno.ProdSequenceNumberRetriever{}
+	return retriever.GetSequenceNumber(name, version)
+}
+
+// SetSequenceNumber sets the sequence number for the given extension name and version.
+// It takes the extension name, extension version, and sequence number as parameters.
+// The sequence number is an integer that represents the order in which the extension was installed.
+// It returns an error if there was a problem setting the sequence number.
+func (s *SeqNumManager) SetSequenceNumber(name, version string, seqNo uint) error {
+	return seqno.SetSequenceNumber(name, version, seqNo)
+}
+
+// FindSeqNum returns the requested the sequence number from either the environment variable or
+// the most recently used file under the config folder.
+// Note that this is different than just choosing the highest number, which may be incorrect
+func (s *SeqNumManager) FindSeqNum(configFolder string) (uint, error) {
+	return seqno.FindSeqNum(logging.NewNopLogger(), configFolder)
+}
+
+// GetCurrentSequenceNumber returns the current sequence number the extension is using
+func (s *SeqNumManager) GetCurrentSequenceNumber(lg log.Logger, name, version string) (sn uint, _ error) {
+	sequenceNumber, err := s.GetSequenceNumber(name, version)
+	if err == extensionerrors.ErrNotFound || err == extensionerrors.ErrNoMrseqFile {
+		// If we can't find the sequence number, then it's possible that the extension
+		// hasn't been installed yet. Go back to 0.
+		lg.Log("event", "Couldn't find current sequence number, likely first execution of the extension, returning sequence number 0")
+		return 0, nil
+	}
+
+	return sequenceNumber, err
+}
+
+func New() SequenceNumberManager {
+	return &SeqNumManager{}
+}

main/cmds.go

@@ -12,8 +12,8 @@ import (
 	"github.com/pkg/errors"
 )
 
-type cmdFunc func(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum int) (msg string, err error)
-type preFunc func(lg log.Logger, seqNum int) error
+type cmdFunc func(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum uint) (msg string, err error)
+type preFunc func(lg log.Logger, seqNum uint) error
 
 type cmd struct {
 	f                  cmdFunc // associated function
@@ -29,7 +29,7 @@ const (
 
 var (
 	cmdInstall   = cmd{install, "Install", false, nil, 52}
-	cmdEnable    = cmd{enable, "Enable", true, nil, 3}
+	cmdEnable    = cmd{enable, "Enable", true, enablePre, 3}
 	cmdUninstall = cmd{uninstall, "Uninstall", false, nil, 3}
 
 	cmds = map[string]cmd{
@@ -41,12 +41,12 @@ var (
 	}
 )
 
-func noop(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum int) (string, error) {
+func noop(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum uint) (string, error) {
 	lg.Log("event", "noop")
 	return "", nil
 }
 
-func install(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum int) (string, error) {
+func install(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum uint) (string, error) {
 	if err := os.MkdirAll(dataDir, 0755); err != nil {
 		return "", errors.Wrap(err, "failed to create data dir")
 	}
@@ -56,7 +56,7 @@ func install(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum int) (strin
 	return "", nil
 }
 
-func uninstall(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum int) (string, error) {
+func uninstall(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum uint) (string, error) {
 	{ // a new context scope with path
 		lg = log.With(lg, "path", dataDir)
 		sendTelemetry(lg, telemetry.EventLevelInfo, telemetry.AppHealthTask, "Removing data dir")
@@ -77,7 +77,29 @@ var (
 	errTerminated = errors.New("Application health process terminated")
 )
 
-func enable(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum int) (string, error) {
+func enablePre(lg log.Logger, seqNum uint) error {
+	// exit if this sequence number (a snapshot of the configuration) is already
+	// processed. if not, save this sequence number before proceeding.
+
+	mrSeqNum, err := seqnoManager.GetCurrentSequenceNumber(lg, fullName, "")
+	if err != nil {
+		return errors.Wrap(err, "failed to get current sequence number")
+	}
+	// If the most recent sequence number is greater than or equal to the requested sequence number,
+	// then the script has already been run and we should exit.
+	if mrSeqNum != 0 && seqNum < mrSeqNum {
+		lg.Log("event", "exit", "message", "the script configuration has already been processed, will not run again")
+		return errors.Errorf("most recent sequence number %d is greater than the requested sequence number %d", mrSeqNum, seqNum)
+	}
+
+	// save the sequence number
+	if err := seqnoManager.SetSequenceNumber(fullName, "", seqNum); err != nil {
+		return errors.Wrap(err, "failed to save sequence number")
+	}
+	return nil
+}
+
+func enable(lg log.Logger, h *handlerenv.HandlerEnvironment, seqNum uint) (string, error) {
 	// parse the extension handler settings (not available prior to 'enable')
 	cfg, err := parseAndValidateSettings(lg, h.ConfigFolder)
 	if err != nil {

main/cmds_test.go

@@ -3,7 +3,11 @@ package main
 import (
 	"testing"
 
+	"github.com/Azure/applicationhealth-extension-linux/internal/seqno"
+	"github.com/go-kit/log"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 )
 
 func Test_commandsExist(t *testing.T) {
@@ -29,3 +33,58 @@ func Test_commands_shouldReportStatus(t *testing.T) {
 	require.True(t, cmds["disable"].shouldReportStatus, "disable should report status")
 	require.True(t, cmds["update"].shouldReportStatus, "update should report status")
 }
+
+func Test_enablePre(t *testing.T) {
+	var (
+		logger          = log.NewNopLogger()
+		seqNumToProcess uint
+		ctrl            = gomock.NewController(t)
+	)
+
+	mockSeqNumManager := seqno.NewMockSequenceNumberManager(ctrl)
+	t.Run("SaveSequenceNumberError_ShouldFail", func(t *testing.T) {
+		// seqNumToProcess = 0, mrSeqNum = 1
+		seqNumToProcess = 0
+		mockSeqNumManager.EXPECT().GetCurrentSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(uint(1), nil)
+		seqnoManager = mockSeqNumManager
+		err := enablePre(logger, seqNumToProcess)
+		assert.Error(t, err)
+		assert.EqualError(t, err, "most recent sequence number 1 is greater than the requested sequence number 0")
+	})
+	t.Run("GetSequenceNumberIsGreaterThanRequestedSequenceNumber_ShouldFail", func(t *testing.T) {
+		// seqNumToProcess = 4, mrSeqNum = 8
+		seqNumToProcess = 4
+		mockSeqNumManager.EXPECT().GetCurrentSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(uint(8), nil)
+		seqnoManager = mockSeqNumManager
+		err := enablePre(logger, seqNumToProcess)
+		assert.Error(t, err)
+		assert.EqualError(t, err, "most recent sequence number 8 is greater than the requested sequence number 4")
+	})
+	t.Run("SequenceNumberisZero_Startup", func(t *testing.T) {
+		// seqNumToProcess = 0, mrSeqNum = 0
+		seqNumToProcess = 0
+		mockSeqNumManager.EXPECT().GetCurrentSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(uint(0), nil)
+		mockSeqNumManager.EXPECT().SetSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)
+		seqnoManager = mockSeqNumManager
+		err := enablePre(logger, seqNumToProcess)
+		assert.NoError(t, err)
+	})
+	t.Run("SequenceNumberAlreadyProcessed", func(t *testing.T) {
+		// seqNumToProcess = 5, mrSeqNum = 5
+		seqNumToProcess = 5
+		seqnoManager = mockSeqNumManager
+		mockSeqNumManager.EXPECT().GetCurrentSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(uint(5), nil)
+		mockSeqNumManager.EXPECT().SetSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)
+		err := enablePre(logger, seqNumToProcess)
+		assert.NoError(t, err)
+	})
+	t.Run("MostRecentSeqNumIsSmaller_ShouldPass", func(t *testing.T) {
+		// seqNumToProcess = 4, mrSeqNum = 2
+		seqNumToProcess = 4
+		mockSeqNumManager.EXPECT().GetCurrentSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(uint(2), nil)
+		mockSeqNumManager.EXPECT().SetSequenceNumber(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil)
+		seqnoManager = mockSeqNumManager
+		err := enablePre(logger, seqNumToProcess)
+		assert.NoError(t, err)
+	})
+}

main/main.go

@@ -9,6 +9,7 @@ import (
 	"syscall"
 
 	"github.com/Azure/applicationhealth-extension-linux/internal/handlerenv"
+	"github.com/Azure/applicationhealth-extension-linux/internal/seqno"
 	"github.com/Azure/applicationhealth-extension-linux/internal/telemetry"
 	"github.com/Azure/applicationhealth-extension-linux/pkg/logging"
 	"github.com/Azure/azure-extension-platform/pkg/extensionevents"
@@ -28,6 +29,8 @@ var (
 	eem *extensionevents.ExtensionEventManager
 
 	sendTelemetry telemetry.LogEventFunc
+
+	seqnoManager seqno.SequenceNumberManager = seqno.New()
 )
 
 func main() {
@@ -61,7 +64,8 @@ func main() {
 		logger.Log("message", "failed to parse handlerenv", "error", err)
 		os.Exit(cmd.failExitCode)
 	}
-	seqNum, err := FindSeqNum(hEnv.ConfigFolder)
+
+	seqNum, err := seqnoManager.FindSeqNum(hEnv.ConfigFolder)
 	if err != nil {
 		logger.Log("message", "failed to find sequence number", "error", err)
 	}

main/reportstatus.go

@@ -14,7 +14,7 @@ import (
 // status.
 //
 // If an error occurs reporting the status, it will be logged and returned.
-func reportStatus(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum int, t StatusType, c cmd, msg string) error {
+func reportStatus(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum uint, t StatusType, c cmd, msg string) error {
 	if !c.shouldReportStatus {
 		lg.Log("status", "not reported for operation (by design)")
 		return nil
@@ -29,7 +29,7 @@ func reportStatus(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum int
 	return nil
 }
 
-func reportStatusWithSubstatuses(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum int, t StatusType, op string, msg string, substatuses []SubstatusItem) error {
+func reportStatusWithSubstatuses(lg log.Logger, hEnv *handlerenv.HandlerEnvironment, seqNum uint, t StatusType, op string, msg string, substatuses []SubstatusItem) error {
 	s := NewStatus(t, op, msg)
 	for _, substatus := range substatuses {
 		s.AddSubstatusItem(substatus)