Use only the latest request to envd (#1293)

Author: jakubno

packages/envd/internal/api/init.go

@@ -9,13 +9,16 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/rs/zerolog"
 	"github.com/txn2/txeh"
 	"golang.org/x/sys/unix"
 
 	"github.com/e2b-dev/infra/packages/envd/internal/host"
 	"github.com/e2b-dev/infra/packages/envd/internal/logs"
 )
 
+var ErrAccessTokenAlreadySet = errors.New("access token is already set")
+
 func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 	defer r.Body.Close()
 
@@ -33,37 +36,23 @@ func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 			return
 		}
 
-		if initRequest.Timestamp != nil {
-			logger.Debug().Msgf("Setting sandbox start time to: %v", *initRequest.Timestamp)
-			ts := unix.NsecToTimespec(initRequest.Timestamp.UnixNano())
-			err = unix.ClockSettime(unix.CLOCK_REALTIME, &ts)
-			if err != nil {
-				logger.Error().Msgf("Failed to set system time: %v", err)
-			}
-		}
-
-		if initRequest.EnvVars != nil {
-			logger.Debug().Msg(fmt.Sprintf("Setting %d env vars", len(*initRequest.EnvVars)))
-
-			for key, value := range *initRequest.EnvVars {
-				logger.Debug().Msgf("Setting env var for %s", key)
-				a.envVars.Store(key, value)
-			}
-		}
+		a.initLock.Lock()
+		defer a.initLock.Unlock()
 
-		if initRequest.AccessToken != nil {
-			if a.accessToken != nil && *initRequest.AccessToken != *a.accessToken {
-				logger.Error().Msg("Access token is already set and cannot be changed")
-				w.WriteHeader(http.StatusConflict)
+		// Update data only if the request is newer or if there's no timestamp at all
+		if initRequest.Timestamp == nil || a.lastSetTime.SetToGreater(initRequest.Timestamp.UnixNano()) {
+			err = a.SetData(logger, initRequest)
+			if err != nil {
+				switch {
+				case errors.Is(err, ErrAccessTokenAlreadySet):
+					w.WriteHeader(http.StatusConflict)
+				default:
+					logger.Error().Msgf("Failed to set data: %v", err)
+					w.WriteHeader(http.StatusBadRequest)
+				}
+				w.Write([]byte(err.Error()))
 				return
 			}
-
-			logger.Debug().Msg("Setting access token")
-			a.accessToken = initRequest.AccessToken
-		}
-
-		if initRequest.HyperloopIP != nil {
-			go a.SetupHyperloop(*initRequest.HyperloopIP)
 		}
 	}
 
@@ -79,6 +68,42 @@ func (a *API) PostInit(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusNoContent)
 }
 
+func (a *API) SetData(logger zerolog.Logger, data PostInitJSONBody) error {
+	if data.Timestamp != nil {
+		logger.Debug().Msgf("Setting sandbox start time to: %v", *data.Timestamp)
+		ts := unix.NsecToTimespec(data.Timestamp.UnixNano())
+		err := unix.ClockSettime(unix.CLOCK_REALTIME, &ts)
+		if err != nil {
+			logger.Error().Msgf("Failed to set system time: %v", err)
+		}
+	}
+
+	if data.EnvVars != nil {
+		logger.Debug().Msg(fmt.Sprintf("Setting %d env vars", len(*data.EnvVars)))
+
+		for key, value := range *data.EnvVars {
+			logger.Debug().Msgf("Setting env var for %s", key)
+			a.envVars.Store(key, value)
+		}
+	}
+
+	if data.AccessToken != nil {
+		if a.accessToken != nil && *data.AccessToken != *a.accessToken {
+			logger.Error().Msg("Access token is already set and cannot be changed")
+			return ErrAccessTokenAlreadySet
+		}
+
+		logger.Debug().Msg("Setting access token")
+		a.accessToken = data.AccessToken
+	}
+
+	if data.HyperloopIP != nil {
+		go a.SetupHyperloop(*data.HyperloopIP)
+	}
+
+	return nil
+}
+
 func (a *API) SetupHyperloop(address string) {
 	a.hyperloopLock.Lock()
 	defer a.hyperloopLock.Unlock()

packages/envd/internal/api/store.go

@@ -18,10 +18,13 @@ type API struct {
 	envVars       *utils.Map[string, string]
 	mmdsChan      chan *host.MMDSOpts
 	hyperloopLock sync.Mutex
+
+	lastSetTime *utils.AtomicMax
+	initLock    sync.Mutex
 }
 
 func New(l *zerolog.Logger, envVars *utils.Map[string, string], mmdsChan chan *host.MMDSOpts, isNotFC bool) *API {
-	return &API{logger: l, envVars: envVars, mmdsChan: mmdsChan, isNotFC: isNotFC}
+	return &API{logger: l, envVars: envVars, mmdsChan: mmdsChan, isNotFC: isNotFC, lastSetTime: utils.NewAtomicMax()}
 }
 
 func (a *API) GetHealth(w http.ResponseWriter, r *http.Request) {

packages/envd/internal/utils/atomic.go

@@ -0,0 +1,26 @@
+package utils
+
+import (
+	"sync"
+)
+
+type AtomicMax struct {
+	val int64
+	mu  sync.Mutex
+}
+
+func NewAtomicMax() *AtomicMax {
+	return &AtomicMax{}
+}
+
+func (a *AtomicMax) SetToGreater(newValue int64) bool {
+	a.mu.Lock()
+	defer a.mu.Unlock()
+
+	if a.val > newValue {
+		return false
+	}
+
+	a.val = newValue
+	return true
+}

packages/envd/internal/utils/atomic_test.go

@@ -0,0 +1,70 @@
+package utils
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAtomicMax_NewAtomicMax(t *testing.T) {
+	am := NewAtomicMax()
+	require.NotNil(t, am)
+	require.Equal(t, int64(0), am.val)
+}
+
+func TestAtomicMax_SetToGreater_InitialValue(t *testing.T) {
+	am := NewAtomicMax()
+
+	// Should succeed when newValue > current
+	assert.True(t, am.SetToGreater(10))
+	assert.Equal(t, int64(10), am.val)
+}
+
+func TestAtomicMax_SetToGreater_EqualValue(t *testing.T) {
+	am := NewAtomicMax()
+	am.val = 10
+
+	// Should succeed when newValue > current
+	assert.True(t, am.SetToGreater(20))
+	assert.Equal(t, int64(20), am.val)
+}
+
+func TestAtomicMax_SetToGreater_GreaterValue(t *testing.T) {
+	am := NewAtomicMax()
+	am.val = 10
+
+	// Should fail when newValue < current, keeping the max value
+	assert.False(t, am.SetToGreater(5))
+	assert.Equal(t, int64(10), am.val)
+}
+
+func TestAtomicMax_SetToGreater_NegativeValues(t *testing.T) {
+	am := NewAtomicMax()
+	am.val = -5
+
+	assert.True(t, am.SetToGreater(-2))
+	assert.Equal(t, int64(-2), am.val)
+}
+
+func TestAtomicMax_SetToGreater_Concurrent(t *testing.T) {
+	am := NewAtomicMax()
+	var wg sync.WaitGroup
+
+	// Run 100 goroutines trying to update the value concurrently
+	numGoroutines := 100
+	wg.Add(numGoroutines)
+
+	for i := range numGoroutines {
+		go func(val int64) {
+			defer wg.Done()
+			am.SetToGreater(val)
+		}(int64(i))
+	}
+
+	wg.Wait()
+
+	// The final value should be 99 (the maximum value)
+	assert.Equal(t, int64(99), am.val)
+}