add DST and week calculation tests

coder · deansheather · Jul 20, 2023 · Jun 20, 2023 · Jun 20, 2023 · Jun 21, 2023
commit eb46ae246f59ce8214d16322dbc7784ec99fd861
diff --git a/coderd/schedule/autostop.go b/coderd/schedule/autostop.go
@@ -118,24 +118,9 @@ func CalculateAutostop(ctx context.Context, params CalculateAutostopParams) (Aut
 			// change the startOfDay to be the Monday of the next applicable
 			// week.
 			if templateSchedule.RestartRequirement.Weeks > 1 {
-				epoch := TemplateRestartRequirementEpoch(loc)
-				if startOfStopDay.Before(epoch) {
-					return autostop, xerrors.New("coder server system clock is incorrect, cannot calculate template restart requirement")
-				}
-				since := startOfStopDay.Sub(epoch)
-				weeksSinceEpoch := int64(since.Hours() / (24 * 7))
-				requiredWeeks := templateSchedule.RestartRequirement.Weeks
-				weeksRemainder := weeksSinceEpoch % requiredWeeks
-				if weeksRemainder != 0 {
-					// Add (requiredWeeks - weeksSince) * 7 days to the current
-					// startOfStopDay, then truncate to Monday midnight.
-					//
-					// This sets startOfStopDay to Monday at midnight of the
-					// next applicable week.
-					y, mo, d := startOfStopDay.Date()
-					d += int(requiredWeeks-weeksRemainder) * 7
-					startOfStopDay = time.Date(y, mo, d, 0, 0, 0, 0, loc)
-					startOfStopDay = truncateMondayMidnight(startOfStopDay)
+				startOfStopDay, err = GetNextApplicableMondayOfNWeeks(startOfStopDay, templateSchedule.RestartRequirement.Weeks)
+				if err != nil {
+					return autostop, xerrors.Errorf("determine start of stop week: %w", err)
 				}
 			}
 
@@ -239,3 +224,102 @@ func truncateMondayMidnight(t time.Time) time.Time {
 	t = time.Date(yy, mm, dd, 0, 0, 0, 0, t.Location())
 	return truncateMidnight(t)
 }
+
+// WeeksSinceEpoch gets the weeks since the epoch for a given time. This is a
+// 0-indexed number of weeks since the epoch (Monday).
+//
+// The timezone embedded in the time object is used to determine the epoch.
+func WeeksSinceEpoch(now time.Time) (int64, error) {
+	epoch := TemplateRestartRequirementEpoch(now.Location())
+	if now.Before(epoch) {
+		return 0, xerrors.New("coder server system clock is incorrect, cannot calculate template restart requirement")
+	}
+
+	// This calculation needs to be done using YearDay, as dividing by the
+	// amount of hours is impacted by daylight savings. Even though daylight
+	// savings is usually only an hour difference, this calculation is used to
+	// get the current week number and could result in an entire week getting
+	// skipped if the calculation is off by an hour.
+	//
+	// Old naive algorithm: weeksSinceEpoch := int64(since.Hours() / (24 * 7))
+
+	// Get days since epoch. Start with a negative number of days, as we want to
+	// subtract the YearDay() of the epoch itself.
+	days := -epoch.YearDay()
+	for i := epoch.Year(); i < now.Year(); i++ {
+		startOfNextYear := time.Date(i+1, 1, 1, 0, 0, 0, 0, now.Location())
+		if startOfNextYear.Year() != i+1 {
+			return 0, xerrors.New("overflow calculating weeks since epoch")
+		}
+		endOfThisYear := startOfNextYear.AddDate(0, 0, -1)
+		if endOfThisYear.Year() != i {
+			return 0, xerrors.New("overflow calculating weeks since epoch")
+		}
+
+		days += endOfThisYear.YearDay()
+	}
+	// Add this year's days.
+	days += now.YearDay()
+
+	// Ensure that the number of days is positive.
+	if days < 0 {
+		return 0, xerrors.New("overflow calculating weeks since epoch")
+	}
+
+	// Divide by 7 to get the number of weeks.
+	weeksSinceEpoch := int64(days / 7)
+	return weeksSinceEpoch, nil
+}
+
+// GetMondayOfWeek gets the Monday (0:00) of the n-th week since epoch.
+func GetMondayOfWeek(loc *time.Location, n int64) (time.Time, error) {
+	if n < 0 {
+		return time.Time{}, xerrors.New("weeks since epoch must be positive")
+	}
+	epoch := TemplateRestartRequirementEpoch(loc)
+	monday := epoch.AddDate(0, 0, int(n*7))
+
+	y, m, d := monday.Date()
+	monday = time.Date(y, m, d, 0, 0, 0, 0, loc)
+	if monday.Weekday() != time.Monday {
+		return time.Time{}, xerrors.Errorf("calculated incorrect Monday for week %v since epoch (actual weekday %q)", n, monday.Weekday())
+	}
+	return monday, nil
+}
+
+// GetNextApplicableMondayOfNWeeks gets the next Monday (0:00) of the next week
+// divisible by n since epoch. If the next applicable week is invalid for any
+// reason, the week after will be used instead (up to 2 attempts).
+//
+// If the current week is divisible by n, then the provided time is returned as
+// is.
+//
+// The timezone embedded in the time object is used to determine the epoch.
+func GetNextApplicableMondayOfNWeeks(now time.Time, n int64) (time.Time, error) {
+	// Get the current week number.
+	weeksSinceEpoch, err := WeeksSinceEpoch(now)
+	if err != nil {
+		return time.Time{}, xerrors.Errorf("get current week number: %w", err)
+	}
+
+	// Get the next week divisible by n.
+	remainder := weeksSinceEpoch % n
+	week := weeksSinceEpoch + (n - remainder)
+	if remainder == 0 {
+		return now, nil
+	}
+
+	// Loop until we find a week that doesn't fail.
+	var lastErr error
+	for i := int64(0); i < 3; i++ {
+		monday, err := GetMondayOfWeek(now.Location(), week+i)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+
+		return monday, nil
+	}
+
+	return time.Time{}, xerrors.Errorf("get next applicable Monday of %v weeks: %w", n, lastErr)
+}
diff --git a/coderd/schedule/autostop_test.go b/coderd/schedule/autostop_test.go
@@ -3,6 +3,7 @@ package schedule_test
 import (
 	"context"
 	"database/sql"
+	"fmt"
 	"testing"
 	"time"
 
@@ -40,6 +41,32 @@ func TestCalculateAutoStop(t *testing.T) {
 	t.Log("fridayEveningSydney", fridayEveningSydney)
 	t.Log("saturdayMidnightSydney", saturdayMidnightSydney)
 
+	dstIn := time.Date(2023, 10, 1, 2, 0, 0, 0, sydneyLoc)   // 1 hour backward
+	dstInQuietHours := "CRON_TZ=Australia/Sydney 30 2 * * *" // never
+	// The expected behavior is that we will pick the next time that falls on
+	// quiet hours after the DST transition. In this case, it will be the same
+	// time the next day.
+	dstInQuietHoursExpectedTime := time.Date(2023, 10, 2, 2, 30, 0, 0, sydneyLoc)
+	beforeDstIn := time.Date(2023, 10, 1, 0, 0, 0, 0, sydneyLoc)
+	saturdayMidnightAfterDstIn := time.Date(2023, 10, 7, 0, 0, 0, 0, sydneyLoc)
+
+	// Wednesday after DST starts.
+	duringDst := time.Date(2023, 10, 4, 0, 0, 0, 0, sydneyLoc)
+	saturdayMidnightAfterDuringDst := saturdayMidnightAfterDstIn
+
+	dstOut := time.Date(2024, 4, 7, 3, 0, 0, 0, sydneyLoc)                        // 1 hour forward
+	dstOutQuietHours := "CRON_TZ=Australia/Sydney 30 3 * * *"                     // twice
+	dstOutQuietHoursExpectedTime := time.Date(2024, 4, 7, 3, 30, 0, 0, sydneyLoc) // in reality, this is the first occurrence
+	beforeDstOut := time.Date(2024, 4, 7, 0, 0, 0, 0, sydneyLoc)
+	saturdayMidnightAfterDstOut := time.Date(2024, 4, 13, 0, 0, 0, 0, sydneyLoc)
+
+	t.Log("dstIn", dstIn)
+	t.Log("beforeDstIn", beforeDstIn)
+	t.Log("saturdayMidnightAfterDstIn", saturdayMidnightAfterDstIn)
+	t.Log("dstOut", dstOut)
+	t.Log("beforeDstOut", beforeDstOut)
+	t.Log("saturdayMidnightAfterDstOut", saturdayMidnightAfterDstOut)
+
 	cases := []struct {
 		name                  string
 		now                   time.Time
@@ -268,6 +295,78 @@ func TestCalculateAutoStop(t *testing.T) {
 			workspaceTTL: 0,
 			errContains:  "coder server system clock is incorrect",
 		},
+		{
+			name:                   "DaylightSavings/OK",
+			now:                    duringDst,
+			templateAllowAutostop:  true,
+			templateDefaultTTL:     0,
+			userQuietHoursSchedule: sydneyQuietHours,
+			templateRestartRequirement: schedule.TemplateRestartRequirement{
+				DaysOfWeek: 0b00100000, // Saturday
+				Weeks:      1,          // weekly
+			},
+			workspaceTTL: 0,
+			// expectedDeadline is copied from expectedMaxDeadline.
+			expectedMaxDeadline: saturdayMidnightAfterDuringDst,
+		},
+		{
+			name:                   "DaylightSavings/SwitchMidWeek/In",
+			now:                    beforeDstIn,
+			templateAllowAutostop:  true,
+			templateDefaultTTL:     0,
+			userQuietHoursSchedule: sydneyQuietHours,
+			templateRestartRequirement: schedule.TemplateRestartRequirement{
+				DaysOfWeek: 0b00100000, // Saturday
+				Weeks:      1,          // weekly
+			},
+			workspaceTTL: 0,
+			// expectedDeadline is copied from expectedMaxDeadline.
+			expectedMaxDeadline: saturdayMidnightAfterDstIn,
+		},
+		{
+			name:                   "DaylightSavings/SwitchMidWeek/Out",
+			now:                    beforeDstOut,
+			templateAllowAutostop:  true,
+			templateDefaultTTL:     0,
+			userQuietHoursSchedule: sydneyQuietHours,
+			templateRestartRequirement: schedule.TemplateRestartRequirement{
+				DaysOfWeek: 0b00100000, // Saturday
+				Weeks:      1,          // weekly
+			},
+			workspaceTTL: 0,
+			// expectedDeadline is copied from expectedMaxDeadline.
+			expectedMaxDeadline: saturdayMidnightAfterDstOut,
+		},
+		{
+			name:                   "DaylightSavings/QuietHoursFallsOnDstSwitch/In",
+			now:                    beforeDstIn.Add(-24 * time.Hour),
+			templateAllowAutostop:  true,
+			templateDefaultTTL:     0,
+			userQuietHoursSchedule: dstInQuietHours,
+			templateRestartRequirement: schedule.TemplateRestartRequirement{
+				DaysOfWeek: 0b01000000, // Sunday
+				Weeks:      1,          // weekly
+			},
+			workspaceTTL: 0,
+			// expectedDeadline is copied from expectedMaxDeadline.
+			expectedMaxDeadline: dstInQuietHoursExpectedTime,
+		},
+		{
+			name:                   "DaylightSavings/QuietHoursFallsOnDstSwitch/Out",
+			now:                    beforeDstOut.Add(-24 * time.Hour),
+			templateAllowAutostop:  true,
+			templateDefaultTTL:     0,
+			userQuietHoursSchedule: dstOutQuietHours,
+			templateRestartRequirement: schedule.TemplateRestartRequirement{
+				DaysOfWeek: 0b01000000, // Sunday
+				Weeks:      1,          // weekly
+			},
+			workspaceTTL: 0,
+			// expectedDeadline is copied from expectedMaxDeadline.
+			expectedMaxDeadline: dstOutQuietHoursExpectedTime,
+		},
+
+		// TODO(@dean): remove max_ttl tests
 		{
 			name:                   "RestartRequirementIgnoresMaxTTL",
 			now:                    fridayEveningSydney.In(time.UTC),
@@ -411,3 +510,74 @@ func TestCalculateAutoStop(t *testing.T) {
 		})
 	}
 }
+
+func TestFindWeek(t *testing.T) {
+	t.Parallel()
+
+	timezones := []string{
+		"UTC",
+		"America/Los_Angeles",
+		"America/New_York",
+		"Europe/Dublin",
+		"Europe/London",
+		"Europe/Paris",
+		"Asia/Kolkata", // India (UTC+5:30)
+		"Asia/Tokyo",
+		"Australia/Sydney",
+		"Australia/Brisbane",
+	}
+
+	for _, tz := range timezones {
+		tz := tz
+		t.Run("Loc/"+tz, func(t *testing.T) {
+			t.Parallel()
+
+			loc, err := time.LoadLocation(tz)
+			require.NoError(t, err)
+
+			now := time.Now().In(loc)
+			currentWeek, err := schedule.WeeksSinceEpoch(now)
+			require.NoError(t, err)
+
+			currentWeekMondayExpected := now.AddDate(0, 0, -int(now.Weekday())+1)
+			y, m, d := currentWeekMondayExpected.Date()
+			currentWeekMondayExpected = time.Date(y, m, d, 0, 0, 0, 0, loc)
+			currentWeekMonday, err := schedule.GetMondayOfWeek(now.Location(), currentWeek)
+			require.NoError(t, err)
+			require.Equal(t, currentWeekMondayExpected, currentWeekMonday)
+
+			t.Log("now", now)
+			t.Log("currentWeek", currentWeek)
+			t.Log("currentMonday", currentWeekMonday)
+
+			// Loop through every single Monday and Sunday for the next 100
+			// years and make sure the week calculations are correct.
+			for i := int64(1); i < 52*100; i++ {
+				msg := fmt.Sprintf("week %d", i)
+
+				monday := currentWeekMonday.AddDate(0, 0, int(i*7))
+				y, m, d := monday.Date()
+				monday = time.Date(y, m, d, 0, 0, 0, 0, loc)
+				require.Equal(t, monday.Weekday(), time.Monday, msg)
+				t.Log(msg, "monday", monday)
+
+				week, err := schedule.WeeksSinceEpoch(monday)
+				require.NoError(t, err, msg)
+				require.Equal(t, currentWeek+i, week, msg)
+
+				gotMonday, err := schedule.GetMondayOfWeek(monday.Location(), week)
+				require.NoError(t, err, msg)
+				require.Equal(t, monday, gotMonday, msg)
+
+				// Check that we get the same week number for late Sunday.
+				sunday := time.Date(y, m, d+6, 23, 59, 59, 0, loc)
+				require.Equal(t, sunday.Weekday(), time.Sunday, msg)
+				t.Log(msg, "sunday", sunday)
+
+				week, err = schedule.WeeksSinceEpoch(sunday)
+				require.NoError(t, err, msg)
+				require.Equal(t, currentWeek+i, week, msg)
+			}
+		})
+	}
+}