go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/lucictx/deadline.go

// Copyright 2020 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 lucictx

import (
	"context"
	"fmt"
	"math"
	"os"
	"os/signal"
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/system/signals"
)

// signalNotify/signalStop are used to mock signal.Notify in tests.
var signalNotify = signal.Notify
var signalStop = signal.Stop

// DefaultGracePeriod is the value of Deadline.grace_period to assume
// if Deadline is entirely missing in LUCI_CONTEXT.
const DefaultGracePeriod = 30 * time.Second

// DeadlineEvent is the type pushed into the cleanup channel returned by
// AdjustDeadline.
type DeadlineEvent int

// The cleanup channel, when it unblocks, will have an infinite supply of one of
// the following event types.
const (
	// ClosureEvent occurs when the context returned by AdjustDeadline is Done.
	// This is the value you'll get from `cleanup` when it's closed.
	ClosureEvent DeadlineEvent = iota

	// InterruptEvent occurs when a SIGTERM/os.Interrupt was handled to unblock
	// the cleanup channel.
	InterruptEvent

	// TimeoutEvent occurs when the cleanup channel was unblocked due to
	// a timeout on deadline-gracePeriod.
	TimeoutEvent
)

func (de DeadlineEvent) String() string {
	switch de {
	case ClosureEvent:
		return "ClosureEvent"
	case InterruptEvent:
		return "InterruptEvent"
	case TimeoutEvent:
		return "TimeoutEvent"
	default:
		panic(fmt.Sprintf("impossible DeadlineEvent %d", de))
	}
}

// earlier returns the earlier of a and b, treating "zero" as "infinity".
func earlier(a, b time.Time) time.Time {
	if a.IsZero() {
		return b
	}
	if b.IsZero() {
		return a
	}
	if a.Before(b) {
		return a
	}
	return b
}

var softDeadlineKey = "holds <-chan DeadlineEvent"

// TrackSoftDeadline returns a context containing a channel for the
// `SoftDeadlineDone` function in this package.
//
// The "soft" deadline is somewhat like context.WithDeadline, except that it
// participates in the LUCI_CONTEXT['deadline'] protocol as well. On hitting the
// soft deadline (or on an external interrupt signal), the SoftDeadlineDone()
// channel will produce a stream of DeadlineEvents. Once this happens, the
// program has LUCI_CONTEXT['deadline']['grace_period'] seconds until
// ctx.Done(). This is meant to give your program time to do cleanup actions
// with a non-canceled Context.
//
// The soft deadline expires based on the earlier of:
//   - LUCI_CONTEXT['deadline']['soft_deadline']
//   - ctx.Deadline() - LUCI_CONTEXT['deadline']['grace_period']
//
// If LUCI_CONTEXT['deadline'] is missing, it is assumed to be:
//
//	{soft_deadline: infinity, grace_period: 30}
//
// This function additionally allows you to reserve a portion of the
// grace_period with `reserveGracePeriod`. This will have the effect of
// adjusting LUCI_CONTEXT['deadline']['grace_period'] in the returned
// context, as well as canceling the returned context that much earlier.
//
// NOTE: If you want to reduce LUCI_CONTEXT['deadline']['soft_deadline'],
// you should do so by applying a Deadline/Timeout to the context prior
// to invoking this function.
//
// Panics if:
//   - reserveGracePeriod < 0.
//   - LUCI_CONTEXT['deadline']['grace_period'] (or its default 30s value)
//     is insufficient to cover reserveGracePeriod.
//
// Example:
//
//	func MainFunc(ctx context.Context) {
//	  // ctx.Deadline  = <unset>
//	  // soft_deadline = t0 + 5:00
//	  // grace_period  = 40
//
//	  newCtx, shutdown := lucictx.TrackSoftDeadline(ctx, 500*time.Millisecond)
//	  defer shutdown()
//	  ScopedFunction(newCtx)
//	}
//
//	func ScopedFunction(newCtx context.Context) {
//	  // hard deadline is (soft_deadline + grace_period - reserveGracePeriod)
//	  // newCtx.Deadline  = unix(t0+5:39.5)
//	  //
//	  // soft_deadline is unchanged
//	  // soft_deadline = t0 + 5:00
//	  //
//	  // grace_period is reduced by reserveGracePeriod
//	  // grace_period = 39.5
//
//	  go func() {
//	    // unblocked at SIGTERM, soft_deadline or shutdown(), whichever is first.
//	    <-lucictx.SoftDeadlineDone()
//	    // have 39.5s to do something (say, send SIGTERM to a child, or start
//	    // tearing down work in-process) before newCtx.Done().
//	  }()
//	}
//
// NOTE: In the event that `ctx` is canceled from outside, `newCtx` will also
// immediately cancel, and SoftDeadlineDone will also move to the ClosureEvent
// state.
func TrackSoftDeadline(ctx context.Context, reserveGracePeriod time.Duration) (newCtx context.Context, shutdown func()) {
	if reserveGracePeriod < 0 {
		panic(errors.Reason("reserveGracePeriod(%d) < 0", reserveGracePeriod).Err())
	}

	d := GetDeadline(ctx)
	if d == nil {
		logging.Warningf(
			ctx, "AdjustDeadline without Deadline in LUCI_CONTEXT. "+
				"Assuming Deadline={grace_period: %.2f}", DefaultGracePeriod.Seconds())
		d = &Deadline{GracePeriod: DefaultGracePeriod.Seconds()}
	}

	needSet := false // set to true if we need to do a write to LUCI_CONTEXT

	// find current soft deadline
	var ctxSoftDeadline time.Time
	if ctxDl, ok := ctx.Deadline(); ok {
		// synthesize the current soft deadline by applying the current grace period
		// to the context's existing hard deadline.
		ctxSoftDeadline = ctxDl.Add(-d.GracePeriodDuration())
	}
	lucictxSoftDeadline := d.SoftDeadlineTime()
	softDeadline := earlier(ctxSoftDeadline, lucictxSoftDeadline)

	// adjust grace period.
	adjustedGrace := d.GracePeriodDuration()
	if reserveGracePeriod > 0 {
		adjustedGrace -= reserveGracePeriod
		if adjustedGrace < 0 {
			panic(errors.Reason(
				"reserveGracePeriod(%s) > gracePeriod(%s)", reserveGracePeriod, d.GracePeriodDuration()).Err())
		}
		d.GracePeriod = adjustedGrace.Seconds()
		needSet = true
	}

	var newCtxCancel func()
	var enforceSoftDeadline bool

	// set the new hard deadline, if any, on newCtx
	if softDeadline.IsZero() /* a.k.a. "infinity" */ {
		// need cancel func here so that newCtx can hit hard closure after
		// a signal/shutdown, even though it won't have a deadline.
		newCtx, newCtxCancel = context.WithCancel(ctx)
	} else {
		if !softDeadline.Equal(lucictxSoftDeadline) {
			d.SetSoftDeadline(softDeadline)
			needSet = true
			enforceSoftDeadline = true
		}
		newCtx, newCtxCancel = clock.WithDeadline(ctx, softDeadline.Add(adjustedGrace))
	}

	if needSet {
		newCtx = SetDeadline(newCtx, d)
	}

	cleanup, shutdown := runDeadlineMonitor(newCtx, newCtxCancel, softDeadline, enforceSoftDeadline, adjustedGrace)
	newCtx = context.WithValue(newCtx, &softDeadlineKey, cleanup)

	return
}

// SoftDeadlineDone is the counterpart of TrackSoftDeadline, and returns
// a channel which unblocks when the soft deadline of ctx is met, or `ctx` has
// been shut down/interrupted.
//
// If ctx does not come from TrackSoftDeadline(), this returns a nil
// channel (i.e. blocks forever)
func SoftDeadlineDone(ctx context.Context) (ret <-chan DeadlineEvent) {
	ret, _ = ctx.Value(&softDeadlineKey).(<-chan DeadlineEvent)
	return
}

func runDeadlineMonitor(ctx context.Context, cancel func(), softDeadline time.Time, enforceSoftDeadline bool, gracePeriod time.Duration) (<-chan DeadlineEvent, func()) {
	cleanupCh := make(chan DeadlineEvent)
	// buffer 1 is essential; otherwise signals will be missed if our goroutine
	// isn't currently blocked on sigCh. With a buffer of 1, sigCh is always ready
	// to send.
	sigCh := make(chan os.Signal, 1)
	signalNotify(sigCh, signals.Interrupts()...)

	var timeoutC <-chan clock.TimerResult
	if enforceSoftDeadline {
		timeoutC = clock.After(ctx, clock.Until(ctx, softDeadline))
	}

	go func() {
		defer cancel()

		evt := func() DeadlineEvent {
			defer signalStop(sigCh)

			select {
			case <-ctx.Done():
				return ClosureEvent

			case <-timeoutC:
				// clock timer channels unblock on ctx.Done; It's a race whether the
				// previous select case or this one will activate, so check to see if
				// the context has Err set, and return accordingly.
				if ctx.Err() != nil {
					return ClosureEvent
				}
				return TimeoutEvent

			case <-sigCh:
				if !softDeadline.IsZero() && clock.Now(ctx).After(softDeadline) {
					return TimeoutEvent
				}
				return InterruptEvent
			}
		}()

		// Note we do this before signaling cleanupCh so that tests can force
		// `clock.After` to run before incrementing the test clock.
		if evt == InterruptEvent {
			timeoutC = clock.After(ctx, gracePeriod)
		} else {
			timeoutC = nil
		}

		if evt == ClosureEvent {
			close(cleanupCh)
		} else {
			go func() {
				for {
					cleanupCh <- evt
				}
			}()
		}

		select {
		case <-timeoutC:
		case <-ctx.Done():
		}
		// note `defer cancel()` at the top; at this point ctx has either timed out
		// from its internal deadline, or we're about to cancel it.
	}()

	return cleanupCh, func() {
		// shutdown func just interrupts on sigCh; multiple calls will have no
		// effect since we only listen to sigCh exactly once.
		select {
		case sigCh <- os.Interrupt:
		default:
		}
	}
}

// GetDeadline retrieves the raw Deadline information from the context.
//
// You probably want to use AdjustDeadline instead.
func GetDeadline(ctx context.Context) *Deadline {
	t := Deadline{}
	ok, err := Lookup(ctx, "deadline", &t)
	if err != nil {
		panic(err)
	}
	if !ok {
		return nil
	}
	return &t
}

// SetDeadline sets the raw Deadline information in the context.
//
// If d is nil, sets a default deadline of:
//
//	{grace_period: DefaultGracePeriod}
//
// If d.deadline == 0, adjusts it to ctx.Deadline() - d.grace_period.
//
// You probably want to use AdjustDeadline instead.
func SetDeadline(ctx context.Context, d *Deadline) context.Context {
	if d == nil {
		d = &Deadline{GracePeriod: DefaultGracePeriod.Seconds()}
	}
	if deadline, ok := ctx.Deadline(); ok && d.SoftDeadline == 0 {
		d.SetSoftDeadline(deadline)
		d.SoftDeadline -= d.GracePeriod
	}
	return Set(ctx, "deadline", d)
}

// SoftDeadlineTime returns the SoftDeadline as a time.Time.
//
// If SoftDeadline is 0 (or *Deadline is nil) this returns a Zero Time.
func (d *Deadline) SoftDeadlineTime() time.Time {
	if d.GetSoftDeadline() == 0 {
		return time.Time{}
	}

	int, frac := math.Modf(d.SoftDeadline)
	return time.Unix(int64(int), int64(frac*1e9)).UTC()
}

// SetSoftDeadline sets the SoftDeadline from a time.Time.
//
// If t.IsZero, this sets the SoftDeadline to 0 as well.
func (d *Deadline) SetSoftDeadline(t time.Time) {
	if t.IsZero() {
		d.SoftDeadline = 0
	} else {
		d.SoftDeadline = float64(t.Unix()) + (float64(t.Nanosecond()) / 1e9)
	}
}

// GracePeriodDuration returns the GracePeriod as a time.Duration.
//
// If d == nil, returns DefaultGracePeriod.
func (d *Deadline) GracePeriodDuration() time.Duration {
	if d == nil {
		return DefaultGracePeriod
	}
	return time.Duration(d.GracePeriod * float64(time.Second))
}