Replication

virtual/activation_cache.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"runtime"
+	"strings"
 	"sync"
 	"time"
 
@@ -83,16 +84,22 @@ func (a *activationsCache) ensureActivation(
 	moduleID,
 	actorID string,
 
-	blacklistedServerID string,
+	extraReplicas uint64,
+	blacklistedServerIDs []string,
 ) ([]types.ActorReference, error) {
 	// Ensure we have a short timeout when communicating with registry.
 	ctx, cc := context.WithTimeout(ctx, defaultActivationCacheTimeout)
 	defer cc()
 
+	isServerIdBlacklisted := make(map[string]bool)
+	for _, s := range blacklistedServerIDs {
+		isServerIdBlacklisted[s] = true
+	}
+
 	if a.c == nil {
 		// Cache disabled, load directly.
 		return a.ensureActivationAndUpdateCache(
-			ctx, namespace, moduleID, actorID, nil, blacklistedServerID)
+			ctx, namespace, moduleID, actorID, extraReplicas, nil, isServerIdBlacklisted, blacklistedServerIDs)
 	}
 
 	var (
@@ -102,18 +109,31 @@ func (a *activationsCache) ensureActivation(
 	bufIface, cacheKey = actorCacheKeyUnsafePooled(namespace, moduleID, actorID)
 	aceI, ok := a.c.Get(cacheKey)
 	bufPool.Put(bufIface)
-	// Cache miss, fill the cache.
-	if !ok ||
+
+	hasBlacklistedID := false
+	if ok {
+		blacklistedIDs := aceI.(activationCacheEntry).blacklistedServerIDs
+
+		for _, id := range blacklistedIDs {
+			if isServerIdBlacklisted[id] {
+				hasBlacklistedID = true
+				break
+			}
+		}
+	}
+
+	// Cache miss or not enough replicas, fill the cache.
+	if !ok || (1+extraReplicas)> uint64(len(aceI.(activationCacheEntry).references)) ||
 		// There is an existing cache entry, however, it was satisfied by a request that did not provide
 		// the same blacklistedServerID we have currently. We must ignore this entry because it could be
 		// stale and end up routing us back to the blacklisted server ID.
-		(blacklistedServerID != "" && aceI.(activationCacheEntry).blacklistedServerID != blacklistedServerID) {
+		hasBlacklistedID {
 		var cachedReferences []types.ActorReference
 		if ok {
 			cachedReferences = aceI.(activationCacheEntry).references
 		}
 		return a.ensureActivationAndUpdateCache(
-			ctx, namespace, moduleID, actorID, cachedReferences, blacklistedServerID)
+			ctx, namespace, moduleID, actorID, extraReplicas, cachedReferences, isServerIdBlacklisted, blacklistedServerIDs)
 	}
 
 	// Cache hit, return result from cache but check if we should proactively refresh
@@ -126,7 +146,7 @@ func (a *activationsCache) ensureActivation(
 		go func() {
 			defer cc()
 			_, err := a.ensureActivationAndUpdateCache(
-				ctx, namespace, moduleID, actorID, ace.references, blacklistedServerID)
+				ctx, namespace, moduleID, actorID, extraReplicas, ace.references, isServerIdBlacklisted, blacklistedServerIDs)
 			if err != nil {
 				a.logger.Error(
 					"error refreshing activation cache in background",
@@ -156,8 +176,10 @@ func (a *activationsCache) ensureActivationAndUpdateCache(
 	moduleID,
 	actorID string,
 
+	extraReplicas uint64,
 	cachedReferences []types.ActorReference,
-	blacklistedServerID string,
+	isServerIdBlacklisted map[string]bool,
+	blacklistedServerIDs []string,
 ) ([]types.ActorReference, error) {
 	// Since this method is less common (cache miss) we just allocate instead of messing
 	// around with unsafe object pooling.
@@ -166,7 +188,7 @@ func (a *activationsCache) ensureActivationAndUpdateCache(
 	// Include blacklistedServerID in the dedupeKey so that "force refreshes" due to a
 	// server blacklist / load-shedding an actor can be initiated *after* a regular
 	// refresh has already started, but *before* it has completed.
-	dedupeKey := fmt.Sprintf("%s::%s", cacheKey, blacklistedServerID)
+	dedupeKey := fmt.Sprintf("%s::%s", cacheKey, strings.Join(blacklistedServerIDs, ":"))
 	referencesI, err, _ := a.deduper.Do(dedupeKey, func() (any, error) {
 		var cachedServerIDs []string
 		for _, ref := range cachedReferences {
@@ -185,7 +207,8 @@ func (a *activationsCache) ensureActivationAndUpdateCache(
 			ModuleID:  moduleID,
 			ActorID:   actorID,
 
-			BlacklistedServerID:       blacklistedServerID,
+			ExtraReplicas:             extraReplicas,
+			BlacklistedServerIDs:      blacklistedServerIDs,
 			CachedActivationServerIDs: cachedServerIDs,
 		})
 		// Release the semaphore as soon as we're done with the network call since the purpose
@@ -209,10 +232,10 @@ func (a *activationsCache) ensureActivationAndUpdateCache(
 		}
 
 		for _, ref := range references.References {
-			if ref.ServerID() == blacklistedServerID {
+			if isServerIdBlacklisted[ref.ServerID()] {
 				return nil, fmt.Errorf(
 					"[invariant violated] registry returned blacklisted server ID: %s in references",
-					blacklistedServerID)
+					blacklistedServerIDs)
 			}
 		}
 
@@ -225,7 +248,7 @@ func (a *activationsCache) ensureActivationAndUpdateCache(
 			references:           references.References,
 			cachedAt:             time.Now(),
 			registryVersionStamp: references.VersionStamp,
-			blacklistedServerID:  blacklistedServerID,
+			blacklistedServerIDs: blacklistedServerIDs,
 		}
 
 		// a.c is internally synchronized, but we use a lock here so we can do an atomic
@@ -264,5 +287,5 @@ type activationCacheEntry struct {
 	references           []types.ActorReference
 	cachedAt             time.Time
 	registryVersionStamp int64
-	blacklistedServerID  string
+	blacklistedServerIDs []string
 }

virtual/activations.go

@@ -118,7 +118,7 @@ func (a *activations) invoke(
 	invokePayload []byte,
 	isTimer bool,
 ) (io.ReadCloser, error) {
-	if err := a.isBlacklisted(reference); err != nil {
+	if err := a.isServerIdBlacklisted(reference); err != nil {
 		return nil, err
 	}
 
@@ -607,7 +607,7 @@ func (a *activations) close(ctx context.Context, numWorkers int) error {
 	return nil
 }
 
-func (a *activations) isBlacklisted(
+func (a *activations) isServerIdBlacklisted(
 	reference types.ActorReferenceVirtual,
 ) error {
 	bufIface, cacheKey := actorCacheKeyUnsafePooled(
@@ -619,7 +619,7 @@ func (a *activations) isBlacklisted(
 		err := fmt.Errorf(
 			"actor %s is blacklisted on this server", reference.ActorID())
 		serverID, _ := a.getServerState()
-		return NewBlacklistedActivationError(err, serverID)
+		return NewBlacklistedActivationError(err, []string{serverID})
 	}
 
 	return nil

virtual/client.go

@@ -45,7 +45,7 @@ func (h *httpClient) InvokeActorRemote(
 
 	req, err := http.NewRequestWithContext(
 		ctx, "POST",
-		fmt.Sprintf("http://%s/api/v1/invoke-actor-direct", reference.Address()),
+		fmt.Sprintf("http://%s/api/v1/invoke-actor-direct", reference.ServerState().Address()),
 		bytes.NewReader(marshaled))
 	if err != nil {
 		return nil, fmt.Errorf("HTTPClient: InvokeDirect: error constructing request: %w", err)
@@ -69,7 +69,7 @@ func (h *httpClient) InvokeActorRemote(
 		// in statusCodeToErrorWrapper will be converted back to the proper in memory
 		// error type if sent by a server to a client.
 		if wrapper, ok := statusCodeToErrorWrapper[resp.StatusCode]; ok {
-			err = wrapper(err, reference.ServerID())
+			err = wrapper(err, []string{reference.ServerID()})
 		}
 		return nil, err
 	}

virtual/environment.go

@@ -9,6 +9,7 @@ import (
 	"io/ioutil"
 	"net"
 	"runtime"
+	"sort"
 	"sync"
 	"time"
 
@@ -411,28 +412,28 @@ func (r *environment) InvokeActorStream(
 	create types.CreateIfNotExist,
 ) (io.ReadCloser, error) {
 	resp, err := r.invokeActorStreamHelper(
-		ctx, namespace, actorID, moduleID, operation, payload, create, "")
+		ctx, namespace, actorID, moduleID, operation, payload, create, []string{})
 	if err == nil {
 		return resp, nil
 	}
 
-	if IsBlacklistedActivationError(err) {
+	if isServerIdBlacklistedActivationError(err) {
 		// If we received an error because the target server has blacklisted activations
 		// of this actor, then we'll invalidate our cache to force the subsequent call
 		// to lookup the actor's new activation location in the registry. We'll also set
 		// a flag on the call to the registry to indicate that the server has blacklisted
 		// that actor to ensure we get an activation on a different serer since the registry
 		// may not know about the blacklist yet.
 		r.activationsCache.delete(namespace, moduleID, actorID)
-		blacklistedServerID := err.(BlacklistedActivationErr).ServerID()
+		blacklistedServerIDs := err.(BlacklistedActivationErr).ServerIDs()
 
 		r.log.Warn(
 			"encountered blacklisted actor, forcing activation cache refresh and retrying",
 			slog.String("actor_id", fmt.Sprintf("%s::%s::%s", namespace, moduleID, actorID)),
-			slog.String("blacklisted_server_id", blacklistedServerID))
+			slog.Any("blacklisted_server_ids", blacklistedServerIDs))
 
 		return r.invokeActorStreamHelper(
-			ctx, namespace, actorID, moduleID, operation, payload, create, blacklistedServerID)
+			ctx, namespace, actorID, moduleID, operation, payload, create, blacklistedServerIDs)
 	}
 
 	return nil, err
@@ -446,7 +447,7 @@ func (r *environment) invokeActorStreamHelper(
 	operation string,
 	payload []byte,
 	create types.CreateIfNotExist,
-	blacklistedServerID string,
+	blacklistedServerID []string,
 ) (io.ReadCloser, error) {
 	if r.isClosed() {
 		return nil, ErrEnvironmentClosed
@@ -468,8 +469,8 @@ func (r *environment) invokeActorStreamHelper(
 	}
 
 	references, err := r.activationsCache.ensureActivation(
-		ctx, namespace, moduleID, actorID, blacklistedServerID)
-	if err != nil {
+		ctx, namespace, moduleID, actorID, create.Options.ExtraReplicas, blacklistedServerID)
+		if err != nil {
 		return nil, fmt.Errorf("error ensuring actor activation: %w", err)
 	}
 	if len(references) == 0 {
@@ -764,14 +765,13 @@ func (r *environment) invokeReferences(
 	payload []byte,
 	create types.CreateIfNotExist,
 ) (io.ReadCloser, error) {
-	// TODO: Load balancing or some other strategy if the number of references is > 1?
-	ref := references[0]
+	ref := pickServerForInvocation(references, create)
 	if !r.opts.ForceRemoteProcedureCalls {
 		// First check the global localEnvironmentsRouter map for scenarios where we're
 		// potentially trying to communicate between multiple different in-memory
 		// instances of Environment.
 		localEnvironmentsRouterLock.RLock()
-		localEnv, ok := localEnvironmentsRouter[ref.Address()]
+		localEnv, ok := localEnvironmentsRouter[ref.ServerState().Address()]
 		localEnvironmentsRouterLock.RUnlock()
 		if ok {
 			return localEnv.InvokeActorDirectStream(
@@ -795,7 +795,7 @@ func (r *environment) invokeReferences(
 		// always return dnsregistry.Localhost as the address for all actor references and
 		// thus ensure that tests can be written without having to also ensure that a NOLA
 		// server is running on the appropriate port, among other things.
-		if ref.Address() == Localhost || ref.Address() == dnsregistry.Localhost {
+		if ref.ServerState().Address() == Localhost || ref.ServerState().Address() == dnsregistry.Localhost {
 			return localEnv.InvokeActorDirectStream(
 				ctx, versionStamp, ref.ServerID(), ref.ServerVersion(), ref,
 				operation, payload, create)
@@ -899,3 +899,12 @@ func formatActorCacheKey(
 	dst = append(dst, []byte(actorID)...)
 	return dst
 }
+
+func pickServerForInvocation(references []types.ActorReference, create types.CreateIfNotExist) types.ActorReference {
+	// TODO: implement invokation strategies in 'create' e.g. memory-balanced, cpu-balanced, multi-invoke...
+	sort.Slice(references, func(i, j int) bool {
+		return references[i].ServerState().UsedMemory() < references[j].ServerState().UsedMemory()
+	})
+
+	return references[0]
+}