go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/appengine/tsmon/tasknum_test.go

// Copyright 2016 The LUCI Authors.
//
// 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 tsmon

import (
	"context"
	"fmt"
	"testing"
	"time"

	"go.chromium.org/luci/appengine/gaetesting"
	"go.chromium.org/luci/common/clock/testclock"
	"go.chromium.org/luci/common/tsmon"
	"go.chromium.org/luci/common/tsmon/target"
	"go.chromium.org/luci/gae/service/datastore"
	srvtsmon "go.chromium.org/luci/server/tsmon"

	. "github.com/smartystreets/goconvey/convey"
)

func TestFindGaps(t *testing.T) {
	t.Parallel()

	tests := []struct {
		numbers        []int
		wantFirst5Gaps []int
	}{
		{[]int{1}, []int{0, 2, 3, 4, 5}},
		{[]int{-1, 1}, []int{0, 2, 3, 4, 5}},
		{[]int{1, 3, 5}, []int{0, 2, 4, 6, 7}},
		{[]int{5, 3, 1}, []int{0, 2, 4, 6, 7}},
		{[]int{3, 1, 5}, []int{0, 2, 4, 6, 7}},
		{[]int{4}, []int{0, 1, 2, 3, 5}},
	}

	for i, test := range tests {
		Convey(fmt.Sprintf("%d. %v", i, test.numbers), t, func() {
			numbers := map[int]struct{}{}
			for _, n := range test.numbers {
				numbers[n] = struct{}{}
			}

			nextNum := gapFinder(numbers)

			// Read 5 numbers from the channel.
			var got []int
			for i := 0; i < 5; i++ {
				got = append(got, nextNum())
			}

			So(got, ShouldResemble, test.wantFirst5Gaps)
		})
	}
}

func buildGAETestContext() (context.Context, testclock.TestClock) {
	c := gaetesting.TestingContext()
	c, clock := testclock.UseTime(c, testclock.TestTimeUTC)
	datastore.GetTestable(c).Consistent(true)
	c, _, _ = tsmon.WithFakes(c)
	return c, clock
}

func TestAssignTaskNumbers(t *testing.T) {
	t.Parallel()

	task1 := target.Task{HostName: "1"}
	task2 := target.Task{HostName: "2"}

	allocator := DatastoreTaskNumAllocator{}

	Convey("Assigns task numbers to unassigned instances", t, func() {
		c, _ := buildGAETestContext()

		// Two new processes belonging to different targets.
		_, err := allocator.NotifyTaskIsAlive(c, &task1, "some.id")
		So(err, ShouldEqual, srvtsmon.ErrNoTaskNumber)
		_, err = allocator.NotifyTaskIsAlive(c, &task2, "some.id")
		So(err, ShouldEqual, srvtsmon.ErrNoTaskNumber)

		So(AssignTaskNumbers(c), ShouldBeNil)

		// Both get 0, since they belong to different targets.
		i, err := allocator.NotifyTaskIsAlive(c, &task1, "some.id")
		So(err, ShouldBeNil)
		So(i, ShouldEqual, 0)
		i, err = allocator.NotifyTaskIsAlive(c, &task2, "some.id")
		So(err, ShouldBeNil)
		So(i, ShouldEqual, 0)

		// Once numbers are assigned, AssignTaskNumbers is noop.
		So(AssignTaskNumbers(c), ShouldBeNil)
	})

	Convey("Doesn't assign the same task number", t, func() {
		c, _ := buildGAETestContext()

		// Two processes belonging to a single target.
		allocator.NotifyTaskIsAlive(c, &task1, "some.task.0")
		allocator.NotifyTaskIsAlive(c, &task1, "some.task.1")

		So(AssignTaskNumbers(c), ShouldBeNil)

		// Get different task numbers.
		i, err := allocator.NotifyTaskIsAlive(c, &task1, "some.task.0")
		So(err, ShouldBeNil)
		So(i, ShouldEqual, 0)
		i, err = allocator.NotifyTaskIsAlive(c, &task1, "some.task.1")
		So(err, ShouldBeNil)
		So(i, ShouldEqual, 1)
	})

	Convey("Expires old instances", t, func() {
		c, clock := buildGAETestContext()

		for _, count := range []int{1, taskQueryBatchSize + 1} {
			Convey(fmt.Sprintf("Count: %d", count), func() {
				// Request a bunch of task numbers.
				for i := 0; i < count; i++ {
					_, err := allocator.NotifyTaskIsAlive(c, &task1, fmt.Sprintf("%d", i))
					So(err, ShouldEqual, srvtsmon.ErrNoTaskNumber)
				}

				// Get all the numbers assigned.
				So(AssignTaskNumbers(c), ShouldBeNil)

				// Yep. Assigned.
				numbers := map[int]struct{}{}
				for i := 0; i < count; i++ {
					num, err := allocator.NotifyTaskIsAlive(c, &task1, fmt.Sprintf("%d", i))
					So(err, ShouldBeNil)
					numbers[num] = struct{}{}
				}
				So(len(numbers), ShouldEqual, count)

				// Move time to make number assignments expire.
				clock.Add(instanceExpirationTimeout + time.Second)
				So(AssignTaskNumbers(c), ShouldBeNil)

				// Yep. Expired.
				for i := 0; i < count; i++ {
					_, err := allocator.NotifyTaskIsAlive(c, &task1, fmt.Sprintf("%d", i))
					So(err, ShouldEqual, srvtsmon.ErrNoTaskNumber)
				}
			})
		}
	})
}