From b0d163c4cdcf1d7cb94a51c62fc3713974f46b49 Mon Sep 17 00:00:00 2001
From: Ramasai Tadepalli <ramasai.tadepalli+3108@gmail.com>
Date: Tue, 29 Aug 2023 19:54:47 -0400
Subject: [PATCH] Make activations cache blacklist TTL configurable

---
 virtual/activations.go | 46 ++++++++++++++++++++++++++----------------
 virtual/environment.go |  7 ++++++-
 2 files changed, 35 insertions(+), 18 deletions(-)

diff --git a/virtual/activations.go b/virtual/activations.go
index ff8a1f4..e7b32d6 100644
--- a/virtual/activations.go
+++ b/virtual/activations.go
@@ -24,8 +24,7 @@ import (
 )
 
 const (
-	// TODO: This should be configurable.
-	activationBlacklistCacheTTL = time.Minute
+	defaultActivationBlacklistCacheTTL = time.Minute
 )
 
 type activations struct {
@@ -55,12 +54,18 @@ type activations struct {
 	}
 
 	// Dependencies.
-	registry      registry.Registry
-	moduleStore   registry.ModuleStore
-	environment   Environment
-	goModules     map[types.NamespacedIDNoType]Module
-	customHostFns map[string]func([]byte) ([]byte, error)
-	gcActorsAfter time.Duration
+	registry                    registry.Registry
+	moduleStore                 registry.ModuleStore
+	environment                 Environment
+	goModules                   map[types.NamespacedIDNoType]Module
+	customHostFns               map[string]func([]byte) ([]byte, error)
+	gcActorsAfter               time.Duration
+	activationBlacklistCacheTTL time.Duration
+}
+
+type activationsOptions struct {
+	gcActorsAfter               time.Duration
+	activationBlacklistCacheTTL time.Duration
 }
 
 func newActivations(
@@ -69,11 +74,17 @@ func newActivations(
 	moduleStore registry.ModuleStore,
 	environment Environment,
 	customHostFns map[string]func([]byte) ([]byte, error),
-	gcActorsAfter time.Duration,
+	opts activationsOptions,
 ) *activations {
+	gcActorsAfter := opts.gcActorsAfter
+	activationBlacklistCacheTTL := opts.activationBlacklistCacheTTL
+
 	if gcActorsAfter < 0 {
 		panic(fmt.Sprintf("[invariant violated] illegal value for gcActorsAfter: %d", gcActorsAfter))
 	}
+	if activationBlacklistCacheTTL <= 0 {
+		activationBlacklistCacheTTL = defaultActivationBlacklistCacheTTL
+	}
 
 	blacklist, err := ristretto.NewCache(&ristretto.Config{
 		NumCounters: maxNumActivationsToCache * 10, // * 10 per the docs.
@@ -93,13 +104,14 @@ func newActivations(
 		_blacklist:            blacklist,
 		_actorResourceTracker: newActorResourceTracker(),
 
-		log:           log.With(slog.String("module", "activations")),
-		registry:      registry,
-		moduleStore:   moduleStore,
-		environment:   environment,
-		goModules:     make(map[types.NamespacedIDNoType]Module),
-		customHostFns: customHostFns,
-		gcActorsAfter: gcActorsAfter,
+		log:                         log.With(slog.String("module", "activations")),
+		registry:                    registry,
+		moduleStore:                 moduleStore,
+		environment:                 environment,
+		goModules:                   make(map[types.NamespacedIDNoType]Module),
+		customHostFns:               customHostFns,
+		gcActorsAfter:               gcActorsAfter,
+		activationBlacklistCacheTTL: activationBlacklistCacheTTL,
 	}
 	a._moduleState.modules = make(map[types.NamespacedID]Module)
 	return a
@@ -538,7 +550,7 @@ func (a *activations) shedMemUsage(numBytes int) {
 	for _, v := range toShed {
 		key := formatActorCacheKey(nil, v.id.Namespace, v.id.Module, v.id.ID)
 		if _, ok := a._blacklist.Get(key); !ok {
-			a._blacklist.SetWithTTL(key, nil, 1, activationBlacklistCacheTTL)
+			a._blacklist.SetWithTTL(key, nil, 1, a.activationBlacklistCacheTTL)
 			a.log.Info(
 				"shedding actor to reduce memory usage",
 				slog.String("actor_id", v.id.String()))
diff --git a/virtual/environment.go b/virtual/environment.go
index 81b2815..38bd576 100644
--- a/virtual/environment.go
+++ b/virtual/environment.go
@@ -91,6 +91,8 @@ const (
 type EnvironmentOptions struct {
 	// ActivationCacheTTL is the TTL of the activation cache.
 	ActivationCacheTTL time.Duration
+	// ActivationBlacklistCacheTTL is the TTL of the activations blacklist cache.
+	ActivationBlacklistCacheTTL time.Duration
 	// DisableActivationCache disables the activation cache.
 	DisableActivationCache bool
 	// Discovery contains the discovery options.
@@ -281,7 +283,10 @@ func NewEnvironment(
 	}
 	env.randState.rng = rand.New(rand.NewSource(time.Now().UnixNano()))
 	activations := newActivations(
-		opts.Logger, reg, moduleStore, env, env.opts.CustomHostFns, opts.GCActorsAfterDurationWithNoInvocations)
+		opts.Logger, reg, moduleStore, env, env.opts.CustomHostFns, activationsOptions{
+			gcActorsAfter:               opts.GCActorsAfterDurationWithNoInvocations,
+			activationBlacklistCacheTTL: opts.ActivationBlacklistCacheTTL,
+		})
 	env.activations = activations
 
 	// Skip confusing log if dnsregistry is being used since it doesn't use the registry-based