Actor state client SetWithTTL

actor/actor.go

@@ -16,6 +16,7 @@ package actor
 import (
 	"context"
 	"sync"
+	"time"
 )
 
 // Client is the interface that should be impl by user's actor client.
@@ -205,8 +206,15 @@ type StateManagerContext interface {
 	Add(ctx context.Context, stateName string, value any) error
 	// Get is to get state store of @stateName with type @reply
 	Get(ctx context.Context, stateName string, reply any) error
-	// Set is to set new state store with @stateName and @value
+	// Set sets a state store with @stateName and @value.
+	// You should always use SetWithTTL unless you also intend to implement your
+	// own state expiration logic. This is to prevent the state store from
+	// growing indefinitely.
 	Set(ctx context.Context, stateName string, value any) error
+	// SetWithTTL sets a state store with @stateName and @value, for the given
+	// TTL. After the TTL has passed, the value will no longer be available with
+	// `Get`. Always preferred over `Set`.
+	SetWithTTL(ctx context.Context, stateName string, value any, ttl time.Duration) error
 	// Remove is to remove state store with @stateName
 	Remove(ctx context.Context, stateName string) error
 	// Contains is to check if state store contains @stateName

actor/manager/container.go

@@ -28,7 +28,7 @@ import (
 // Deprecated: use ActorContainerContext instead.
 type ActorContainer interface {
 	Invoke(methodName string, param []byte) ([]reflect.Value, actorErr.ActorErr)
-	//nolint:staticcheck
+	//nolint:staticcheck // SA1019 Deprecated: use ActorContainerContext instead.
 	GetActor() actor.Server
 }
 

actor/state/actor_state_change.go

@@ -13,12 +13,22 @@ limitations under the License.
 
 package state
 
+import (
+	"time"
+)
+
 type ActorStateChange struct {
-	stateName  string
-	value      interface{}
-	changeKind ChangeKind
+	stateName    string
+	value        interface{}
+	changeKind   ChangeKind
+	ttlInSeconds *int64
 }
 
-func NewActorStateChange(stateName string, value interface{}, changeKind ChangeKind) *ActorStateChange {
-	return &ActorStateChange{stateName: stateName, value: value, changeKind: changeKind}
+func NewActorStateChange(stateName string, value any, changeKind ChangeKind, ttl *time.Duration) *ActorStateChange {
+	var ttlF *int64
+	if ttl != nil && *ttl > 0 {
+		ttlInSeconds := int64(ttl.Seconds())
+		ttlF = &ttlInSeconds
+	}
+	return &ActorStateChange{stateName: stateName, value: value, changeKind: changeKind, ttlInSeconds: ttlF}
 }

actor/state/actor_state_change_test.go

@@ -14,36 +14,40 @@ limitations under the License.
 package state
 
 import (
-	"reflect"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestNewActorStateChange(t *testing.T) {
-	type args struct {
+	secs5 := int64(5)
+
+	tests := map[string]struct {
 		stateName  string
-		value      interface{}
+		value      any
 		changeKind ChangeKind
-	}
-	tests := []struct {
-		name string
-		args args
-		want *ActorStateChange
+		ttl        time.Duration
+		want       *ActorStateChange
 	}{
-		{
-			name: "init",
-			args: args{
-				stateName:  "testStateName",
-				value:      "testValue",
-				changeKind: Add,
-			},
-			want: &ActorStateChange{stateName: "testStateName", value: "testValue", changeKind: Add},
+		"init": {
+			stateName:  "testStateName",
+			value:      "testValue",
+			changeKind: Add,
+			ttl:        time.Second*5 + time.Millisecond*400,
+			want:       &ActorStateChange{stateName: "testStateName", value: "testValue", changeKind: Add, ttlInSeconds: &secs5},
+		},
+		"no TTL": {
+			stateName:  "testStateName",
+			value:      "testValue",
+			changeKind: Add,
+			ttl:        0,
+			want:       &ActorStateChange{stateName: "testStateName", value: "testValue", changeKind: Add, ttlInSeconds: nil},
 		},
 	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			if got := NewActorStateChange(tt.args.stateName, tt.args.value, tt.args.changeKind); !reflect.DeepEqual(got, tt.want) {
-				t.Errorf("NewActorStateChange() = %v, want %v", got, tt.want)
-			}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			assert.Equal(t, test.want, NewActorStateChange(test.stateName, test.value, test.changeKind, &test.ttl))
 		})
 	}
 }

actor/state/state_async_provider.go

@@ -100,6 +100,7 @@ func (d *DaprStateAsyncProvider) ApplyContext(ctx context.Context, actorType, ac
 			OperationType: daprOperationName,
 			Key:           stateChange.stateName,
 			Value:         value,
+			TTLInSeconds:  stateChange.ttlInSeconds,
 		})
 	}
 

actor/state/state_change_metadata.go

@@ -13,6 +13,8 @@ limitations under the License.
 
 package state
 
+import "time"
+
 type ChangeKind string
 
 const (
@@ -24,12 +26,18 @@ const (
 
 type ChangeMetadata struct {
 	Kind  ChangeKind
-	Value interface{}
+	Value any
+	TTL   *time.Duration
 }
 
-func NewChangeMetadata(kind ChangeKind, value interface{}) *ChangeMetadata {
+func NewChangeMetadata(kind ChangeKind, value any) *ChangeMetadata {
 	return &ChangeMetadata{
 		Kind:  kind,
 		Value: value,
 	}
 }
+
+func (c *ChangeMetadata) WithTTL(ttl time.Duration) *ChangeMetadata {
+	c.TTL = &ttl
+	return c
+}

actor/state/state_manager.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"reflect"
 	"sync"
+	"time"
 
 	"github.com/dapr/go-sdk/actor"
 )
@@ -152,6 +153,30 @@ func (s *stateManagerCtx) Set(_ context.Context, stateName string, value any) er
 	return nil
 }
 
+func (s *stateManagerCtx) SetWithTTL(_ context.Context, stateName string, value any, ttl time.Duration) error {
+	if stateName == "" {
+		return errors.New("state name can't be empty")
+	}
+
+	if ttl < 0 {
+		return errors.New("ttl can't be negative")
+	}
+
+	if val, ok := s.stateChangeTracker.Load(stateName); ok {
+		metadata := val.(*ChangeMetadata)
+		if metadata.Kind == None || metadata.Kind == Remove {
+			metadata.Kind = Update
+		}
+		s.stateChangeTracker.Store(stateName, NewChangeMetadata(metadata.Kind, value))
+		return nil
+	}
+	s.stateChangeTracker.Store(stateName, (&ChangeMetadata{
+		Kind:  Add,
+		Value: value,
+	}).WithTTL(ttl))
+	return nil
+}
+
 func (s *stateManagerCtx) Remove(ctx context.Context, stateName string) error {
 	if stateName == "" {
 		return errors.New("state name can't be empty")
@@ -200,7 +225,7 @@ func (s *stateManagerCtx) Save(ctx context.Context) error {
 	s.stateChangeTracker.Range(func(key, value any) bool {
 		stateName := key.(string)
 		metadata := value.(*ChangeMetadata)
-		changes = append(changes, NewActorStateChange(stateName, metadata.Value, metadata.Kind))
+		changes = append(changes, NewActorStateChange(stateName, metadata.Value, metadata.Kind, metadata.TTL))
 		return true
 	})
 	if err := s.stateAsyncProvider.ApplyContext(ctx, s.actorTypeName, s.actorID, changes); err != nil {

client/actor.go

@@ -19,6 +19,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"strconv"
 
 	anypb "github.com/golang/protobuf/ptypes/any"
 
@@ -28,6 +29,10 @@ import (
 	pb "github.com/dapr/go-sdk/dapr/proto/runtime/v1"
 )
 
+const (
+	metadataKeyTTLInSeconds = "ttlInSeconds"
+)
+
 type InvokeActorRequest struct {
 	ActorType string
 	ActorID   string
@@ -463,6 +468,7 @@ type ActorStateOperation struct {
 	OperationType string
 	Key           string
 	Value         []byte
+	TTLInSeconds  *int64
 }
 
 func (c *GRPCClient) SaveStateTransactionally(ctx context.Context, actorType, actorID string, operations []*ActorStateOperation) error {
@@ -477,12 +483,18 @@ func (c *GRPCClient) SaveStateTransactionally(ctx context.Context, actorType, ac
 	}
 	grpcOperations := make([]*pb.TransactionalActorStateOperation, 0)
 	for _, op := range operations {
+		var metadata map[string]string
+		if op.TTLInSeconds != nil {
+			metadata = make(map[string]string)
+			metadata[metadataKeyTTLInSeconds] = strconv.FormatInt(*op.TTLInSeconds, 10)
+		}
 		grpcOperations = append(grpcOperations, &pb.TransactionalActorStateOperation{
 			OperationType: op.OperationType,
 			Key:           op.Key,
 			Value: &anypb.Any{
 				Value: op.Value,
 			},
+			Metadata: metadata,
 		})
 	}
 	_, err := c.protoClient.ExecuteActorStateTransaction(c.withAuthToken(ctx), &pb.ExecuteActorStateTransactionRequest{