Merge pull request #2565 from keboola/hosan/PAT-1630

Refactor AppInfo for context support and data retrieval simplification

Author: hosekpeter

internal/pkg/service/appsproxy/dataapps/k8sapp/watcher.go

@@ -120,7 +120,7 @@ func NewStateWatcher(d dependencies, client dynamic.Interface, namespace string)
 
 // GetState returns the cached AppInfo for the app. Returns (AppInfo{}, false) if not yet cached.
 // If the E2B access token is missing but a secret name is known, it attempts to load the token lazily.
-func (w *StateWatcher) GetState(appID api.AppID) (AppInfo, bool) {
+func (w *StateWatcher) GetState(ctx context.Context, appID api.AppID) (AppInfo, bool) {
 	v, ok := w.apps.Load(appID)
 	if !ok {
 		return AppInfo{}, false
@@ -130,15 +130,16 @@ func (w *StateWatcher) GetState(appID api.AppID) (AppInfo, bool) {
 	// Lazy-load E2B token: the Secret may not have existed when the App CRD event was processed.
 	// a singleflight coalesces concurrent requests for the same secret into a single K8s API call.
 	if e.e2bAccessToken == "" && e.e2bSecretName != "" {
+		fetchCtx := context.WithoutCancel(ctx)
 		token, err, _ := w.tokenLoadGroup.Do(e.e2bSecretName, func() (any, error) {
-			return w.loadSecretToken(context.Background(), e.e2bSecretName)
+			return w.loadSecretToken(fetchCtx, e.e2bSecretName)
 		})
 		if err != nil {
-			w.logger.Warnf(context.Background(), "App %s: failed to lazy-load E2B access token from secret %q: %s", appID, e.e2bSecretName, err)
+			w.logger.Warnf(ctx, "App %s: failed to lazy-load E2B access token from secret %q: %s", appID, e.e2bSecretName, err)
 		} else if t, ok := token.(string); t != "" && ok {
 			e.e2bAccessToken = t
 			w.apps.Store(appID, e)
-			w.logger.Infof(context.Background(), "App %s: lazy-loaded E2B access token from secret %q", appID, e.e2bSecretName)
+			w.logger.Infof(ctx, "App %s: lazy-loaded E2B access token from secret %q", appID, e.e2bSecretName)
 		}
 	}
 

internal/pkg/service/appsproxy/dataapps/k8sapp/watcher_test.go

@@ -86,7 +86,7 @@ func TestStateWatcher_GetState_UnknownWhenEmpty(t *testing.T) {
 	fakeClient := newFakeClient()
 	watcher := k8sapp.NewStateWatcher(newTestDeps(t), fakeClient, testNamespace)
 
-	info, ok := watcher.GetState(api.AppID("app-123"))
+	info, ok := watcher.GetState(t.Context(), api.AppID("app-123"))
 	assert.False(t, ok)
 	assert.Empty(t, info.ActualState)
 }
@@ -107,7 +107,7 @@ func TestStateWatcher_GetState_AfterCacheSync(t *testing.T) {
 	watcher := k8sapp.NewStateWatcher(d, fakeClient, testNamespace)
 
 	assert.Eventually(t, func() bool {
-		info, ok := watcher.GetState(api.AppID("app-123"))
+		info, ok := watcher.GetState(t.Context(), api.AppID("app-123"))
 		return ok && info.ActualState == k8sapp.AppActualStateStopped
 	}, 5*time.Second, 50*time.Millisecond)
 }
@@ -128,7 +128,7 @@ func TestStateWatcher_WakeupApp(t *testing.T) {
 
 	// Wait for the informer to cache the object.
 	require.Eventually(t, func() bool {
-		_, ok := watcher.GetState(api.AppID("app-123"))
+		_, ok := watcher.GetState(t.Context(), api.AppID("app-123"))
 		return ok
 	}, 5*time.Second, 50*time.Millisecond)
 
@@ -180,7 +180,7 @@ func TestStateWatcher_GetState_UpstreamTarget(t *testing.T) {
 	var info k8sapp.AppInfo
 	assert.Eventually(t, func() bool {
 		var ok bool
-		info, ok = watcher.GetState(api.AppID("app-123"))
+		info, ok = watcher.GetState(t.Context(), api.AppID("app-123"))
 		return ok && info.UpstreamTarget != nil
 	}, 5*time.Second, 50*time.Millisecond)
 
@@ -205,11 +205,11 @@ func TestStateWatcher_GetState_UpstreamTarget_AbsentWhenMissing(t *testing.T) {
 	watcher := k8sapp.NewStateWatcher(d, fakeClient, testNamespace)
 
 	assert.Eventually(t, func() bool {
-		_, ok := watcher.GetState(api.AppID("app-123"))
+		_, ok := watcher.GetState(t.Context(), api.AppID("app-123"))
 		return ok
 	}, 5*time.Second, 50*time.Millisecond)
 
-	info, ok := watcher.GetState(api.AppID("app-123"))
+	info, ok := watcher.GetState(t.Context(), api.AppID("app-123"))
 	require.True(t, ok)
 	assert.Nil(t, info.UpstreamTarget)
 }
@@ -271,7 +271,7 @@ func TestStateWatcher_GetState_E2BAccessToken(t *testing.T) {
 	var info k8sapp.AppInfo
 	assert.Eventually(t, func() bool {
 		var ok bool
-		info, ok = watcher.GetState(api.AppID("app-e2b"))
+		info, ok = watcher.GetState(t.Context(), api.AppID("app-e2b"))
 		return ok && info.E2BAccessToken != ""
 	}, 5*time.Second, 50*time.Millisecond)
 
@@ -294,11 +294,11 @@ func TestStateWatcher_GetState_E2BAccessToken_MissingSecret(t *testing.T) {
 	watcher := k8sapp.NewStateWatcher(d, fakeClient, testNamespace)
 
 	assert.Eventually(t, func() bool {
-		_, ok := watcher.GetState(api.AppID("app-e2b"))
+		_, ok := watcher.GetState(t.Context(), api.AppID("app-e2b"))
 		return ok
 	}, 5*time.Second, 50*time.Millisecond)
 
-	info, ok := watcher.GetState(api.AppID("app-e2b"))
+	info, ok := watcher.GetState(t.Context(), api.AppID("app-e2b"))
 	require.True(t, ok)
 	assert.Empty(t, info.E2BAccessToken)
 }
@@ -319,11 +319,11 @@ func TestStateWatcher_GetState_NonE2BApp_NoToken(t *testing.T) {
 	watcher := k8sapp.NewStateWatcher(d, fakeClient, testNamespace)
 
 	assert.Eventually(t, func() bool {
-		_, ok := watcher.GetState(api.AppID("app-regular"))
+		_, ok := watcher.GetState(t.Context(), api.AppID("app-regular"))
 		return ok
 	}, 5*time.Second, 50*time.Millisecond)
 
-	info, ok := watcher.GetState(api.AppID("app-regular"))
+	info, ok := watcher.GetState(t.Context(), api.AppID("app-regular"))
 	require.True(t, ok)
 	assert.Empty(t, info.E2BAccessToken)
 }

internal/pkg/service/appsproxy/dataapps/wakeup/wakeup_test.go

@@ -74,7 +74,7 @@ func TestManager_Wakeup(t *testing.T) {
 
 	// Wait for watcher cache to sync.
 	require.Eventually(t, func() bool {
-		_, ok := watcher.GetState(appID)
+		_, ok := watcher.GetState(ctx, appID)
 		return ok
 	}, 5*time.Second, 50*time.Millisecond)
 
@@ -122,7 +122,7 @@ func TestManager_Wakeup_Race(t *testing.T) {
 
 	// Wait for watcher cache to sync.
 	require.Eventually(t, func() bool {
-		_, ok := watcher.GetState(appID)
+		_, ok := watcher.GetState(ctx, appID)
 		return ok
 	}, 5*time.Second, 50*time.Millisecond)
 

internal/pkg/service/appsproxy/proxy/apphandler/manager.go

@@ -2,12 +2,15 @@ package apphandler
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"net/http"
 	"sync"
 
 	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/config"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/dataapps/api"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/dataapps/appconfig"
+	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/dataapps/k8sapp"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/proxy/apphandler/authproxy"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/proxy/apphandler/upstream"
 	"github.com/keboola/keboola-as-code/internal/pkg/service/appsproxy/proxy/pagewriter"
@@ -29,11 +32,10 @@ type Manager struct {
 }
 
 type appHandlerWrapper struct {
-	lock           *sync.Mutex
-	handler        http.Handler
-	cancel         context.CancelCauseFunc
-	upstreamURL    string // appsProxy.upstreamUrl in effect when this handler was created
-	e2bAccessToken string // E2B access token in effect when this handler was created
+	lock        *sync.Mutex
+	handler     http.Handler
+	cancel      context.CancelCauseFunc
+	handlerHash string // hash of UpstreamTarget + E2BAccessToken; handler is recreated when it changes
 }
 
 type dependencies interface {
@@ -72,16 +74,15 @@ func (m *Manager) HandlerFor(ctx context.Context, result appconfig.AppConfigResu
 		return m.newErrorHandler(ctx, api.AppConfig{ID: result.AppID}, result.Err)
 	}
 
-	// Create a new handler, if needed (config changed, upstreamUrl changed, or E2B token changed)
-	currentURL := m.upstreamManager.UpstreamURL(result.AppID)
-	currentToken := m.upstreamManager.E2BAccessToken(result.AppID)
-	if wrapper.handler == nil || result.Modified || wrapper.upstreamURL != currentURL || wrapper.e2bAccessToken != currentToken {
+	// Create a new handler when config changed, upstream URL changed, or E2B token changed.
+	// Only a hash is stored so raw secrets don't linger in the wrapper.
+	currentHash := handlerHash(m.upstreamManager.AppInfo(ctx, result.AppID))
+	if wrapper.handler == nil || result.Modified || wrapper.handlerHash != currentHash {
 		if wrapper.cancel != nil {
 			wrapper.cancel(errors.New("configuration changed"))
 		}
 		wrapper.handler, wrapper.cancel = m.newHandler(ctx, result.AppConfig)
-		wrapper.upstreamURL = currentURL
-		wrapper.e2bAccessToken = currentToken
+		wrapper.handlerHash = currentHash
 	}
 
 	return wrapper.handler
@@ -110,6 +111,21 @@ func (m *Manager) newHandler(ctx context.Context, app api.AppConfig) (http.Handl
 	return handler, appUpstream.Cancel
 }
 
+// handlerHash computes a hash from the fields that require handler recreation when they change.
+// Only a hash is stored so raw secrets don't linger in the wrapper.
+func handlerHash(info k8sapp.AppInfo, ok bool) string {
+	if !ok {
+		return ""
+	}
+	h := sha256.New()
+	if info.UpstreamTarget != nil {
+		h.Write([]byte(info.UpstreamTarget.String()))
+	}
+	h.Write([]byte{0})
+	h.Write([]byte(info.E2BAccessToken))
+	return hex.EncodeToString(h.Sum(nil))
+}
+
 func (m *Manager) newErrorHandler(ctx context.Context, app api.AppConfig, err error) http.Handler {
 	err = svcErrors.WrapWithExceptionID(middleware.RequestIDFromContext(ctx), err)
 	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {

internal/pkg/service/appsproxy/proxy/apphandler/upstream/upstream.go

@@ -98,24 +98,10 @@ func (m *Manager) Shutdown(ctx context.Context) {
 	m.wg.Wait()
 }
 
-// UpstreamURL returns the appsProxy.upstreamUrl string for appID from the K8s cache.
-// Returns "" when the app is not cached or the field is absent/invalid.
-func (m *Manager) UpstreamURL(appID api.AppID) string {
-	info, ok := m.stateWatcher.GetState(appID)
-	if !ok || info.UpstreamTarget == nil {
-		return ""
-	}
-	return info.UpstreamTarget.String()
-}
-
-// E2BAccessToken returns the E2B access token for appID from the K8s cache.
-// Returns "" when the app is not cached or is not an E2B sandbox.
-func (m *Manager) E2BAccessToken(appID api.AppID) string {
-	info, ok := m.stateWatcher.GetState(appID)
-	if !ok {
-		return ""
-	}
-	return info.E2BAccessToken
+// AppInfo returns the cached AppInfo for appID from the K8s cache in a single call.
+// Returns (AppInfo{}, false) when the app is not cached.
+func (m *Manager) AppInfo(ctx context.Context, appID api.AppID) (k8sapp.AppInfo, bool) {
+	return m.stateWatcher.GetState(ctx, appID)
 }
 
 func (m *Manager) NewUpstream(ctx context.Context, app api.AppConfig) (upstream *AppUpstream, err error) {
@@ -124,7 +110,7 @@ func (m *Manager) NewUpstream(ctx context.Context, app api.AppConfig) (upstream
 
 	// Resolve target URL at creation time; immutable after this point.
 	var target *url.URL
-	if info, ok := m.stateWatcher.GetState(app.ID); ok {
+	if info, ok := m.stateWatcher.GetState(ctx, app.ID); ok {
 		target = info.UpstreamTarget // pre-parsed by watcher on CRD event; may be nil
 	}
 
@@ -160,7 +146,7 @@ func (u *AppUpstream) ServeHTTPOrError(rw http.ResponseWriter, req *http.Request
 
 	// K8s state pre-check: if we know the app is not running, handle it synchronously
 	// without attempting DNS/upstream. Falls through if state is unknown or Running.
-	if appInfo, ok := u.manager.stateWatcher.GetState(u.app.ID); ok && appInfo.ActualState != k8sapp.AppActualStateRunning {
+	if appInfo, ok := u.manager.stateWatcher.GetState(ctx, u.app.ID); ok && appInfo.ActualState != k8sapp.AppActualStateRunning {
 		switch {
 		case appInfo.ActualState == k8sapp.AppActualStateStarting:
 			u.manager.pageWriter.WriteSpinnerPage(rw, req, u.app)
@@ -219,8 +205,8 @@ func (u *AppUpstream) newReverseProxy() *httputil.ReverseProxy {
 			// Always fetch the latest token from the state watcher to handle
 			// secret recreation (updates propagate asynchronously).
 			r.Out.Header.Del("e2b-traffic-access-token")
-			if token := u.manager.E2BAccessToken(u.app.ID); token != "" {
-				r.Out.Header.Set("e2b-traffic-access-token", token)
+			if info, ok := u.manager.AppInfo(r.Out.Context(), u.app.ID); ok && info.E2BAccessToken != "" {
+				r.Out.Header.Set("e2b-traffic-access-token", info.E2BAccessToken)
 			}
 		},
 		Transport:    u.manager.transport,

internal/pkg/service/appsproxy/proxy/proxy_test.go

@@ -1862,7 +1862,7 @@ func TestAppProxyRouter(t *testing.T) {
 				)
 				require.NoError(t, err)
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("123"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("123"))
 					return ok && info.ActualState == k8sapp.AppActualStateStopped
 				}, 5*time.Second, 50*time.Millisecond)
 			},
@@ -1881,7 +1881,7 @@ func TestAppProxyRouter(t *testing.T) {
 				)
 				require.NoError(t, err)
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("123"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("123"))
 					return ok && info.ActualState == k8sapp.AppActualStateRunning
 				}, 5*time.Second, 50*time.Millisecond)
 
@@ -1908,7 +1908,7 @@ func TestAppProxyRouter(t *testing.T) {
 				)
 				require.NoError(t, err)
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("123"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("123"))
 					return ok && info.ActualState == k8sapp.AppActualStateStopped
 				}, 5*time.Second, 50*time.Millisecond)
 			},
@@ -1952,7 +1952,7 @@ func TestAppProxyRouter(t *testing.T) {
 				)
 				require.NoError(t, err)
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("oidc"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("oidc"))
 					return ok && info.ActualState == k8sapp.AppActualStateStopped
 				}, 5*time.Second, 50*time.Millisecond)
 			},
@@ -2017,7 +2017,7 @@ func TestAppProxyRouter(t *testing.T) {
 				)
 				require.NoError(t, err)
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("oidc"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("oidc"))
 					return ok && info.ActualState == k8sapp.AppActualStateRunning
 				}, 5*time.Second, 50*time.Millisecond)
 
@@ -2041,7 +2041,7 @@ func TestAppProxyRouter(t *testing.T) {
 				)
 				require.NoError(t, err)
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("oidc"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("oidc"))
 					return ok && info.ActualState == k8sapp.AppActualStateStopped
 				}, 5*time.Second, 50*time.Millisecond)
 			},
@@ -2380,7 +2380,7 @@ func TestAppProxyRouter(t *testing.T) {
 				require.NoError(t, err)
 
 				require.Eventually(t, func() bool {
-					info, ok := watcher.GetState(api.AppID("123"))
+					info, ok := watcher.GetState(t.Context(), api.AppID("123"))
 					return ok && !info.AutoRestartEnabled
 				}, 5*time.Second, 50*time.Millisecond)
 			},