Conformance Tests: State- ttlExpireTime

state/etcd/etcd.go

@@ -192,38 +192,29 @@ func (e *Etcd) doSet(ctx context.Context, key string, val any, etag *string, ttl
 		return err
 	}
 
+	var leaseID clientv3.LeaseID
 	if ttlInSeconds != nil {
-		resp, err := e.client.Grant(ctx, *ttlInSeconds)
+		var resp *clientv3.LeaseGrantResponse
+		resp, err = e.client.Grant(ctx, *ttlInSeconds)
 		if err != nil {
 			return fmt.Errorf("couldn't grant lease %s: %w", key, err)
 		}
-		if etag != nil {
-			etag, _ := strconv.ParseInt(*etag, 10, 64)
-			_, err = e.client.Txn(ctx).
-				If(clientv3.Compare(clientv3.ModRevision(key), "=", etag)).
-				Then(clientv3.OpPut(key, reqVal, clientv3.WithLease(resp.ID))).
-				Commit()
-		} else {
-			_, err = e.client.Put(ctx, key, reqVal, clientv3.WithLease(resp.ID))
-		}
-		if err != nil {
-			return fmt.Errorf("couldn't set key %s: %w", key, err)
-		}
+		leaseID = resp.ID
+	}
+
+	if etag != nil {
+		etag, _ := strconv.ParseInt(*etag, 10, 64)
+		_, err = e.client.Txn(ctx).
+			If(clientv3.Compare(clientv3.ModRevision(key), "=", etag)).
+			Then(clientv3.OpPut(key, reqVal, clientv3.WithLease(leaseID))).
+			Commit()
 	} else {
-		var err error
-		if etag != nil {
-			etag, _ := strconv.ParseInt(*etag, 10, 64)
-			_, err = e.client.Txn(ctx).
-				If(clientv3.Compare(clientv3.ModRevision(key), "=", etag)).
-				Then(clientv3.OpPut(key, reqVal)).
-				Commit()
-		} else {
-			_, err = e.client.Put(ctx, key, reqVal)
-		}
-		if err != nil {
-			return fmt.Errorf("couldn't set key %s: %w", key, err)
-		}
+		_, err = e.client.Put(ctx, key, reqVal, clientv3.WithLease(leaseID))
+	}
+	if err != nil {
+		return fmt.Errorf("couldn't set key %s: %w", key, err)
 	}
+
 	return nil
 }
 

tests/conformance/state/state.go

@@ -990,6 +990,16 @@ func ConformanceTests(t *testing.T, props map[string]string, statestore state.St
 	}
 
 	if config.HasOperation("ttl") {
+		t.Run("set ttl with bad value should error", func(t *testing.T) {
+			require.Error(t, statestore.Set(context.Background(), &state.SetRequest{
+				Key:   key + "-ttl",
+				Value: "⏱️",
+				Metadata: map[string]string{
+					"ttlInSeconds": "foo",
+				},
+			}))
+		})
+
 		t.Run("set and get with TTL", func(t *testing.T) {
 			// Check if ttl feature is listed
 			features := statestore.Features()
@@ -1031,6 +1041,183 @@ func ConformanceTests(t *testing.T, props map[string]string, statestore state.St
 			features := statestore.Features()
 			require.False(t, state.FeatureTTL.IsPresent(features))
 		})
+
+		t.Run("no TTL should not return any expire time", func(t *testing.T) {
+			err := statestore.Set(context.Background(), &state.SetRequest{
+				Key:      key + "-no-ttl",
+				Value:    "⏱️",
+				Metadata: map[string]string{},
+			})
+			require.NoError(t, err)
+
+			// Request immediately
+			res, err := statestore.Get(context.Background(), &state.GetRequest{Key: key + "-no-ttl"})
+			require.NoError(t, err)
+			assertEquals(t, "⏱️", res)
+
+			assert.NotContains(t, res.Metadata, "ttlExpireTime")
+		})
+
+		t.Run("ttlExpireTime", func(t *testing.T) {
+			if !config.HasOperation("transaction") {
+				// This test is only for state stores that support transactions
+				return
+			}
+
+			unsupported := []string{
+				"redis.v6",
+				"redis.v7",
+				"etcd.v1",
+			}
+
+			for _, noSup := range unsupported {
+				if strings.Contains(config.ComponentName, noSup) {
+					t.Skipf("skipping test for unsupported state store %s", noSup)
+				}
+			}
+
+			t.Run("set and get expire time", func(t *testing.T) {
+				now := time.Now()
+				err := statestore.Set(context.Background(), &state.SetRequest{
+					Key:   key + "-ttl-expire-time",
+					Value: "⏱️",
+					Metadata: map[string]string{
+						// Expire in an hour.
+						"ttlInSeconds": "3600",
+					},
+				})
+				require.NoError(t, err)
+
+				// Request immediately
+				res, err := statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+
+				require.Containsf(t, res.Metadata, "ttlExpireTime", "expected metadata to contain ttlExpireTime")
+				expireTime, err := time.Parse(time.RFC3339, res.Metadata["ttlExpireTime"])
+				require.NoError(t, err)
+				assert.InDelta(t, now.Add(time.Hour).UnixMilli(), expireTime.UnixMilli(), float64(time.Minute*10))
+			})
+
+			t.Run("ttl set to -1 should remove the TTL of a state store key", func(t *testing.T) {
+				req := func(meta map[string]string) *state.SetRequest {
+					return &state.SetRequest{
+						Key:      key + "-ttl-expire-time-minus-1",
+						Value:    "⏱️",
+						Metadata: meta,
+					}
+				}
+
+				require.NoError(t, statestore.Set(context.Background(), req(map[string]string{
+					// Expire in 2 seconds.
+					"ttlInSeconds": "2",
+				})))
+
+				// Request immediately
+				res, err := statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time-minus-1",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+				assert.Contains(t, res.Metadata, "ttlExpireTime")
+
+				// Remove TTL by setting a value of -1.
+				require.NoError(t, statestore.Set(context.Background(), req(map[string]string{
+					"ttlInSeconds": "-1",
+				})))
+				res, err = statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time-minus-1",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+				assert.NotContains(t, res.Metadata, "ttlExpireTime")
+
+				// Ensure that the key is not expired after previous TTL.
+				time.Sleep(3 * time.Second)
+
+				res, err = statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time-minus-1",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+
+				// Set a new TTL.
+				require.NoError(t, statestore.Set(context.Background(), req(map[string]string{
+					"ttlInSeconds": "2",
+				})))
+				res, err = statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time-minus-1",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+				assert.Contains(t, res.Metadata, "ttlExpireTime")
+
+				// Remove TTL by omitting the ttlInSeconds field.
+				require.NoError(t, statestore.Set(context.Background(), req(map[string]string{})))
+				res, err = statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time-minus-1",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+				assert.NotContains(t, res.Metadata, "ttlExpireTime")
+
+				// Ensure key is not expired after previous TTL.
+				time.Sleep(3 * time.Second)
+				res, err = statestore.Get(context.Background(), &state.GetRequest{
+					Key: key + "-ttl-expire-time-minus-1",
+				})
+				require.NoError(t, err)
+				assertEquals(t, "⏱️", res)
+				assert.NotContains(t, res.Metadata, "ttlExpireTime")
+			})
+
+			t.Run("set and get expire time bulkGet", func(t *testing.T) {
+				now := time.Now()
+				require.NoError(t, statestore.Set(context.Background(), &state.SetRequest{
+					Key:      key + "-ttl-expire-time-bulk-1",
+					Value:    "123",
+					Metadata: map[string]string{"ttlInSeconds": "3600"},
+				}))
+
+				require.NoError(t, statestore.Set(context.Background(), &state.SetRequest{
+					Key:      key + "-ttl-expire-time-bulk-2",
+					Value:    "234",
+					Metadata: map[string]string{"ttlInSeconds": "3600"},
+				}))
+
+				// Request immediately
+				res, err := statestore.BulkGet(context.Background(), []state.GetRequest{
+					{Key: key + "-ttl-expire-time-bulk-1"},
+					{Key: key + "-ttl-expire-time-bulk-2"},
+				}, state.BulkGetOpts{})
+				require.NoError(t, err)
+
+				require.Len(t, res, 2)
+				sort.Slice(res, func(i, j int) bool {
+					return res[i].Key < res[j].Key
+				})
+
+				assert.Equal(t, key+"-ttl-expire-time-bulk-1", res[0].Key)
+				assert.Equal(t, key+"-ttl-expire-time-bulk-2", res[1].Key)
+				assert.Equal(t, []byte(`"123"`), res[0].Data)
+				assert.Equal(t, []byte(`"234"`), res[1].Data)
+
+				for i := range res {
+					if config.HasOperation("transaction") {
+						require.Containsf(t, res[i].Metadata, "ttlExpireTime", "expected metadata to contain ttlExpireTime")
+						expireTime, err := time.Parse(time.RFC3339, res[i].Metadata["ttlExpireTime"])
+						require.NoError(t, err)
+						// Check the expire time is returned and is in a 10 minute window. This
+						// window should be _more_ than enough.
+						assert.InDelta(t, now.Add(time.Hour).UnixMilli(), expireTime.UnixMilli(), float64(time.Minute*10))
+					} else {
+						assert.NotContains(t, res[i].Metadata, "ttlExpireTime")
+					}
+				}
+			})
+		})
 	}
 }