From fc6ef32f8ae4a978b201d9d7ae2658febe917490 Mon Sep 17 00:00:00 2001
From: Ben Hoyt <ben.hoyt@canonical.com>
Date: Fri, 9 Aug 2024 16:47:17 +1200
Subject: [PATCH] fix(cmdstate): don't store execSetup and environment in state

Per #411, the environment (and other task data) can get quite large for
exec objects. Given that exec commands are one-shot and only relevant
for the current run of Pebble, there's no need to persist them.

Note that none of this data is in "api-data" so it's not accessible via
the API at all right now, and so this is a non-breaking change.

Later if we want to make exec tasks a bit more introspectable we can
add the command line or other fields to "api-data" (but we shouldn't
add the entire environment).

Fixes #411.
---
 internals/overlord/cmdstate/handlers.go |  8 ++++----
 internals/overlord/cmdstate/manager.go  | 16 ++++++++++++++++
 internals/overlord/cmdstate/request.go  |  2 +-
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/internals/overlord/cmdstate/handlers.go b/internals/overlord/cmdstate/handlers.go
index 4e657ef6a..5889e6466 100644
--- a/internals/overlord/cmdstate/handlers.go
+++ b/internals/overlord/cmdstate/handlers.go
@@ -70,13 +70,13 @@ type execution struct {
 }
 
 func (m *CommandManager) doExec(task *state.Task, tomb *tomb.Tomb) error {
-	var setup execSetup
 	st := task.State()
 	st.Lock()
-	err := task.Get("exec-setup", &setup)
+	setupObj := st.Cached(execSetupKey{task.ID()})
 	st.Unlock()
-	if err != nil {
-		return fmt.Errorf("cannot get exec setup object for task %q: %v", task.ID(), err)
+	setup, ok := setupObj.(*execSetup)
+	if !ok || setup == nil {
+		return fmt.Errorf("internal error: cannot get exec setup object for task %q", task.ID())
 	}
 
 	// Set up the object that will track the execution.
diff --git a/internals/overlord/cmdstate/manager.go b/internals/overlord/cmdstate/manager.go
index 2fdd06b46..974687467 100644
--- a/internals/overlord/cmdstate/manager.go
+++ b/internals/overlord/cmdstate/manager.go
@@ -18,6 +18,8 @@ import (
 	"net/http"
 	"sync"
 
+	"gopkg.in/tomb.v2"
+
 	"github.com/canonical/pebble/internals/overlord/state"
 )
 
@@ -34,9 +36,23 @@ func NewManager(runner *state.TaskRunner) *CommandManager {
 		executionsCond: sync.NewCond(&sync.Mutex{}),
 	}
 	runner.AddHandler("exec", manager.doExec, nil)
+
+	// Delete the in-memory execSetup object when the exec is done.
+	runner.AddCleanup("exec", func(task *state.Task, tomb *tomb.Tomb) error {
+		st := task.State()
+		st.Lock()
+		defer st.Unlock()
+		st.Cache(execSetupKey{task.ID()}, nil)
+		return nil
+	})
+
 	return manager
 }
 
+type execSetupKey struct {
+	taskID string
+}
+
 // Ensure is part of the overlord.StateManager interface.
 func (m *CommandManager) Ensure() error {
 	return nil
diff --git a/internals/overlord/cmdstate/request.go b/internals/overlord/cmdstate/request.go
index d14106078..34a0515ce 100644
--- a/internals/overlord/cmdstate/request.go
+++ b/internals/overlord/cmdstate/request.go
@@ -138,7 +138,7 @@ func Exec(st *state.State, args *ExecArgs) (*state.Task, ExecMetadata, error) {
 		GroupID:     args.GroupID,
 		WorkingDir:  workingDir,
 	}
-	task.Set("exec-setup", &setup)
+	st.Cache(execSetupKey{task.ID()}, &setup)
 
 	metadata := ExecMetadata{
 		TaskID:      task.ID(),