github.com/matrixorigin/matrixone@v1.2.0/pkg/taskservice/task_service_cron_test.go

// Copyright 2022 Matrix Origin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package taskservice

import (
	"context"
	"testing"
	"time"

	"github.com/matrixorigin/matrixone/pkg/common/moerr"
	"github.com/matrixorigin/matrixone/pkg/common/runtime"
	"github.com/matrixorigin/matrixone/pkg/common/stopper"
	"github.com/matrixorigin/matrixone/pkg/pb/task"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestScheduleCronTask(t *testing.T) {
	runScheduleCronTaskTest(t, func(store *memTaskStorage, s *taskService, ctx context.Context) {
		assert.NoError(t, s.CreateCronTask(ctx, newTestTaskMetadata("t1"), "*/1 * * * * *"))

		s.StartScheduleCronTask()
		defer s.StopScheduleCronTask()

		waitHasTasks(t, store, time.Second*20, WithTaskParentTaskIDCond(EQ, "t1"))
	})
}

func TestRetryScheduleCronTask(t *testing.T) {
	runScheduleCronTaskTest(t, func(store *memTaskStorage, s *taskService, ctx context.Context) {
		n := 0
		store.preUpdateCron = func() error {
			if n == 0 {
				n++
				return moerr.NewInfo(context.TODO(), "test error")
			}
			return nil
		}

		assert.NoError(t, s.CreateCronTask(ctx, newTestTaskMetadata("t1"), "*/1 * * * * *"))

		s.StartScheduleCronTask()
		defer s.StopScheduleCronTask()

		waitHasTasks(t, store, time.Second*20, WithTaskParentTaskIDCond(EQ, "t1"))
	})
}

func TestScheduleCronTaskImmediately(t *testing.T) {
	runScheduleCronTaskTest(t, func(store *memTaskStorage, s *taskService, ctx context.Context) {
		task := newTestCronTask("t1", "*/1 * * * * *")
		task.CreateAt = time.Now().UnixMilli()
		task.NextTime = task.CreateAt
		task.TriggerTimes = 0
		task.UpdateAt = time.Now().UnixMilli()

		mustAddTestCronTask(t, store, 1, task)

		s.StartScheduleCronTask()
		defer s.StopScheduleCronTask()

		waitHasTasks(t, store, time.Second*20, WithTaskParentTaskIDCond(EQ, "t1"))
	})
}

func TestScheduleCronTaskLimitConcurrency(t *testing.T) {
	runScheduleCronTaskTest(t, func(store *memTaskStorage, s *taskService, ctx context.Context) {
		cronTask := newTestCronTask("t1", "* * * * ? *")
		cronTask.CreateAt = time.Now().UnixMilli()
		cronTask.NextTime = cronTask.CreateAt
		cronTask.TriggerTimes = 0
		cronTask.UpdateAt = time.Now().UnixMilli()
		cronTask.Metadata.Options.Concurrency = 1

		mustAddTestCronTask(t, store, 1, cronTask)

		s.StartScheduleCronTask()
		defer s.StopScheduleCronTask()

		waitHasTasks(t, store, time.Second*20,
			WithTaskParentTaskIDCond(EQ, "t1"))
		assertTaskCountEqual(t, store, time.Second*5, 1,
			WithTaskParentTaskIDCond(EQ, "t1"),
			WithTaskStatusCond(task.TaskStatus_Running))
	})
}

func TestRemovedCronTask(t *testing.T) {
	runScheduleCronTaskTest(t, func(store *memTaskStorage, s *taskService, ctx context.Context) {
		assert.NoError(t, s.CreateCronTask(ctx, newTestTaskMetadata("t1"), "*/1 * * * * *"))

		s.StartScheduleCronTask()
		defer s.StopScheduleCronTask()
		time.Sleep(time.Second * 3)
		s.crons.stopForTest()
		waitJobsCount(t, 1, s, time.Second*10)

		store.Lock()
		store.cronTaskIndexes = make(map[string]uint64)
		store.cronTasks = make(map[uint64]task.CronTask)
		store.Unlock()

		s.crons.startForTest(t, s.fetchCronTasks)
		time.Sleep(time.Second * 3)
		s.crons.stopForTest()
		waitJobsCount(t, 0, s, time.Second*10)
	})
}

func TestReplaceCronTask(t *testing.T) {
	runScheduleCronTaskTest(t, func(store *memTaskStorage, s *taskService, ctx context.Context) {
		assert.NoError(t, s.CreateCronTask(ctx, newTestTaskMetadata("t1"), "*/1 * * * * *"))
		s.StartScheduleCronTask()
		defer s.StopScheduleCronTask()
		time.Sleep(time.Second * 3)
		s.crons.stopForTest()

		jobInCron := s.crons.jobs[1]
		taskInStore := store.cronTasks[jobInCron.task.ID]

		require.Equal(t, jobInCron.task.TriggerTimes, taskInStore.TriggerTimes)

		t.Log("set trigger times to 0")
		jobInCron.task.TriggerTimes = 0
		require.Equal(t, s.crons.jobs[1].task.TriggerTimes, uint64(0))

		s.crons.startForTest(t, s.fetchCronTasks)
		time.Sleep(time.Second * 3)
		s.crons.stopForTest()

		jobInCron = s.crons.jobs[1]
		taskInStore = store.cronTasks[jobInCron.task.ID]
		require.Equal(t, jobInCron.task.TriggerTimes, taskInStore.TriggerTimes)
	})
}

func (c *crons) stopForTest() {
	c.stopper.Stop()
	<-c.cron.Stop().Done()
}

func (c *crons) startForTest(t *testing.T, fn func(ctx context.Context)) {
	c.cron.Start()
	c.stopper = stopper.NewStopper("cronTasks")
	require.NoError(t, c.stopper.RunTask(fn))
}

func runScheduleCronTaskTest(t *testing.T, testFunc func(*memTaskStorage, *taskService, context.Context)) {
	fetchInterval = 300 * time.Millisecond

	store := NewMemTaskStorage().(*memTaskStorage)
	s := NewTaskService(runtime.DefaultRuntime(), store).(*taskService)
	defer func() {
		assert.NoError(t, s.Close())
	}()

	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*10)
	defer cancel()
	testFunc(store, s, ctx)
}

func waitJobsCount(t *testing.T, n int, s *taskService, timeout time.Duration) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	defer func() {
		require.Equal(t, n, len(s.crons.entries))
	}()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			if len(s.crons.entries) == n {
				return
			}
		}
		time.Sleep(time.Millisecond * 10)
	}
}

func waitHasTasks(t *testing.T, store *memTaskStorage, timeout time.Duration, conds ...Condition) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	for {
		select {
		case <-ctx.Done():
			require.Fail(t, "wait any tasks failed")
			return
		default:
			tasks, err := store.QueryAsyncTask(ctx, conds...)
			require.NoError(t, err)
			if len(tasks) > 0 {
				return
			}
		}
		time.Sleep(time.Millisecond * 10)
	}
}

func assertTaskCountEqual(t *testing.T, store *memTaskStorage, timeout time.Duration, count int, conds ...Condition) {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	for {
		select {
		case <-ctx.Done():
			return
		default:
			tasks, err := store.QueryAsyncTask(ctx, conds...)
			require.NoError(t, err)
			require.LessOrEqual(t, len(tasks), count, "err")
		}
		time.Sleep(time.Second)
	}
}