Use Cassandra's ttl to populate Expiry field in persistence.TaskInfo (#6624)

* Use Cassandra's ttl to populate Expiry field in persistence.TaskInfo

Author: fimanishi

common/persistence/nosql/nosqlplugin/cassandra/tasks.go

@@ -351,7 +351,7 @@ func (db *cdb) GetTasksCount(ctx context.Context, filter *nosqlplugin.TasksFilte
 
 // SelectTasks return tasks that associated to a tasklist
 func (db *cdb) SelectTasks(ctx context.Context, filter *nosqlplugin.TasksFilter) ([]*nosqlplugin.TaskRow, error) {
-	// Reading tasklist tasks need to be quorum level consistent, otherwise we could loose task
+	// Reading tasklist tasks need to be quorum level consistent, otherwise we could lose tasks
 	query := db.session.Query(templateGetTasksQuery,
 		filter.DomainID,
 		filter.TaskListName,
@@ -374,8 +374,25 @@ PopulateTasks:
 		if !ok { // no tasks, but static column record returned
 			continue
 		}
+
+		// Extract the TTL value
+		ttlValue, ttlExists := task["ttl"]
+
+		// Check if TTL is null or an integer
+		var ttl *int
+		if ttlExists && ttlValue != nil {
+			if ttlInt, ok := ttlValue.(int); ok {
+				ttl = &ttlInt // TTL is an integer
+			}
+		}
+
 		t := createTaskInfo(task["task"].(map[string]interface{}))
 		t.TaskID = taskID.(int64)
+
+		if ttl != nil {
+			t.Expiry = db.timeSrc.Now().Add(time.Duration(*ttl) * time.Second)
+		}
+
 		response = append(response, t)
 		if len(response) == filter.BatchSize {
 			break PopulateTasks

common/persistence/nosql/nosqlplugin/cassandra/tasks_cql.go

@@ -49,7 +49,7 @@ const (
 		`domain_id, task_list_name, task_list_type, type, task_id, task) ` +
 		`VALUES(?, ?, ?, ?, ?, ` + templateTaskType + `) USING TTL ?`
 
-	templateGetTasksQuery = `SELECT task_id, task ` +
+	templateGetTasksQuery = `SELECT task_id, task, TTL(task) AS ttl ` +
 		`FROM tasks ` +
 		`WHERE domain_id = ? ` +
 		`and task_list_name = ? ` +

common/persistence/nosql/nosqlplugin/cassandra/tasks_test.go

@@ -922,23 +922,25 @@ func TestSelectTasks(t *testing.T) {
 						"task_id": int64(1),
 						"task": map[string]interface{}{
 							"domain_id":        &fakeUUID{uuid: "domain1"},
-							"wofklow_id":       "wid1",
+							"workflow_id":      "wid1",
 							"schedule_id":      int64(42),
 							"created_time":     ts,
 							"run_id":           &fakeUUID{uuid: "runid1"},
 							"partition_config": map[string]string{},
 						},
+						"ttl": 2,
 					},
 					{
 						"task_id": int64(2),
 						"task": map[string]interface{}{
 							"domain_id":        &fakeUUID{uuid: "domain1"},
-							"worklow_id":       "wid1",
+							"workflow_id":      "wid1",
 							"schedule_id":      int64(45),
 							"created_time":     ts,
 							"run_id":           &fakeUUID{uuid: "runid1"},
 							"partition_config": map[string]string{},
 						},
+						"ttl": nil,
 					},
 					{
 						"missing_task_id": int64(1), // missing task_id so this row will be skipped
@@ -947,54 +949,62 @@ func TestSelectTasks(t *testing.T) {
 						"task_id": int64(3),
 						"task": map[string]interface{}{
 							"domain_id":        &fakeUUID{uuid: "domain1"},
-							"worklow_id":       "wid1",
+							"workflow_id":      "wid1",
 							"schedule_id":      int64(48),
 							"created_time":     ts,
 							"run_id":           &fakeUUID{uuid: "runid1"},
 							"partition_config": map[string]string{},
 						},
+						"ttl": 4,
 					},
 					{
 						"task_id": int64(4), // this will be skipped because filter.BatchSize is reached
 						"task": map[string]interface{}{
 							"domain_id":        &fakeUUID{uuid: "domain1"},
-							"worklow_id":       "wid1",
+							"workflow_id":      "wid1",
 							"schedule_id":      int64(51),
 							"created_time":     ts,
 							"run_id":           &fakeUUID{uuid: "runid1"},
 							"partition_config": map[string]string{},
 						},
+						"ttl": 5,
 					},
 				},
 			},
 			wantTasks: []*nosqlplugin.TaskRow{
 				{
 					DomainID:        "domain1",
 					TaskID:          1,
+					WorkflowID:      "wid1",
 					RunID:           "runid1",
 					ScheduledID:     42,
+					Expiry:          ts.Add(time.Second * 2),
 					CreatedTime:     ts,
 					PartitionConfig: map[string]string{},
 				},
 				{
 					DomainID:        "domain1",
 					TaskID:          2,
+					WorkflowID:      "wid1",
 					RunID:           "runid1",
 					ScheduledID:     45,
+					Expiry:          time.Time{},
 					CreatedTime:     ts,
 					PartitionConfig: map[string]string{},
 				},
 				{
 					DomainID:        "domain1",
 					TaskID:          3,
+					WorkflowID:      "wid1",
 					RunID:           "runid1",
 					ScheduledID:     48,
+					Expiry:          ts.Add(time.Second * 4),
 					CreatedTime:     ts,
 					PartitionConfig: map[string]string{},
 				},
 			},
 			wantQueries: []string{
-				`SELECT task_id, task FROM tasks WHERE domain_id = domain1 and task_list_name = tasklist1 and task_list_type = 1 and type = 0 and task_id > 0 and task_id <= 100`,
+				`SELECT task_id, task, TTL(task) AS ttl FROM tasks WHERE domain_id = domain1 and task_list_name = tasklist1 and task_list_type = 1 and type = 0 and task_id > 0 and task_id <= 100`,
 			},
 		},
 	}
@@ -1017,7 +1027,9 @@ func TestSelectTasks(t *testing.T) {
 			client := gocql.NewMockClient(ctrl)
 			cfg := &config.NoSQL{}
 			logger := testlogger.New(t)
-			db := newCassandraDBFromSession(cfg, session, logger, nil, dbWithClient(client))
+			timeSrc := clock.NewMockedTimeSourceAt(ts)
+
+			db := newCassandraDBFromSession(cfg, session, logger, nil, dbWithClient(client), dbWithTimeSource(timeSrc))
 
 			gotRows, err := db.SelectTasks(context.Background(), tc.filter)
 

common/persistence/nosql/nosqlplugin/types.go

@@ -168,6 +168,7 @@ type (
 		WorkflowID      string
 		RunID           string
 		ScheduledID     int64
+		Expiry          time.Time
 		CreatedTime     time.Time
 		PartitionConfig map[string]string
 	}

service/matching/tasklist/task_reader.go

@@ -299,7 +299,7 @@ func (tr *taskReader) getTaskBatch(readLevel, maxReadLevel int64) ([]*persistenc
 }
 
 func (tr *taskReader) isTaskExpired(t *persistence.TaskInfo) bool {
-	return t.Expiry.After(epochStartTime) && tr.timeSource.Now().After(t.Expiry)
+	return !t.Expiry.IsZero() && t.Expiry.After(epochStartTime) && tr.timeSource.Now().After(t.Expiry)
 }
 
 func (tr *taskReader) addTasksToBuffer(tasks []*persistence.TaskInfo) bool {