github.com/opentofu/opentofu@v1.7.1/internal/backend/local/hook_state.go

// Copyright (c) The OpenTofu Authors
// SPDX-License-Identifier: MPL-2.0
// Copyright (c) 2023 HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0

package local

import (
	"log"
	"sync"
	"time"

	"github.com/opentofu/opentofu/internal/states"
	"github.com/opentofu/opentofu/internal/states/statemgr"
	"github.com/opentofu/opentofu/internal/tofu"
)

// StateHook is a hook that continuously updates the state by calling
// WriteState on a statemgr.Full.
type StateHook struct {
	tofu.NilHook
	sync.Mutex

	StateMgr statemgr.Writer

	// If PersistInterval is nonzero then for any new state update after
	// the duration has elapsed we'll try to persist a state snapshot
	// to the persistent backend too.
	// That's only possible if field Schemas is valid, because the
	// StateMgr.PersistState function for some backends needs schemas.
	PersistInterval time.Duration

	// Schemas are the schemas to use when persisting state due to
	// PersistInterval. This is ignored if PersistInterval is zero,
	// and PersistInterval is ignored if this is nil.
	Schemas *tofu.Schemas

	intermediatePersist IntermediateStatePersistInfo
}

type IntermediateStatePersistInfo struct {
	// RequestedPersistInterval is the persist interval requested by whatever
	// instantiated the StateHook.
	//
	// Implementations of [IntermediateStateConditionalPersister] should ideally
	// respect this, but may ignore it if they use something other than the
	// passage of time to make their decision.
	RequestedPersistInterval time.Duration

	// LastPersist is the time when the last intermediate state snapshot was
	// persisted, or the time of the first report for OpenTofu Core if there
	// hasn't yet been a persisted snapshot.
	LastPersist time.Time

	// ForcePersist is true when OpenTofu CLI has received an interrupt
	// signal and is therefore trying to create snapshots more aggressively
	// in anticipation of possibly being terminated ungracefully.
	// [IntermediateStateConditionalPersister] implementations should ideally
	// persist every snapshot they get when this flag is set, unless they have
	// some external information that implies this shouldn't be necessary.
	ForcePersist bool
}

var _ tofu.Hook = (*StateHook)(nil)

func (h *StateHook) PostStateUpdate(new *states.State) (tofu.HookAction, error) {
	h.Lock()
	defer h.Unlock()

	h.intermediatePersist.RequestedPersistInterval = h.PersistInterval

	if h.intermediatePersist.LastPersist.IsZero() {
		// The first PostStateUpdate starts the clock for intermediate
		// calls to PersistState.
		h.intermediatePersist.LastPersist = time.Now()
	}

	if h.StateMgr != nil {
		if err := h.StateMgr.WriteState(new); err != nil {
			return tofu.HookActionHalt, err
		}
		if mgrPersist, ok := h.StateMgr.(statemgr.Persister); ok && h.PersistInterval != 0 && h.Schemas != nil {
			if h.shouldPersist() {
				err := mgrPersist.PersistState(h.Schemas)
				if err != nil {
					return tofu.HookActionHalt, err
				}
				h.intermediatePersist.LastPersist = time.Now()
			} else {
				log.Printf("[DEBUG] State storage %T declined to persist a state snapshot", h.StateMgr)
			}
		}
	}

	return tofu.HookActionContinue, nil
}

func (h *StateHook) Stopping() {
	h.Lock()
	defer h.Unlock()

	// If OpenTofu has been asked to stop then that might mean that a hard
	// kill signal will follow shortly in case OpenTofu doesn't stop
	// quickly enough, and so we'll try to persist the latest state
	// snapshot in the hope that it'll give the user less recovery work to
	// do if they _do_ subsequently hard-kill OpenTofu during an apply.

	if mgrPersist, ok := h.StateMgr.(statemgr.Persister); ok && h.Schemas != nil {
		// While we're in the stopping phase we'll try to persist every
		// new state update to maximize every opportunity we get to avoid
		// losing track of objects that have been created or updated.
		// OpenTofu Core won't start any new operations after it's been
		// stopped, so at most we should see one more PostStateUpdate
		// call per already-active request.
		h.intermediatePersist.ForcePersist = true

		if h.shouldPersist() {
			err := mgrPersist.PersistState(h.Schemas)
			if err != nil {
				// This hook can't affect OpenTofu Core's ongoing behavior,
				// but it's a best effort thing anyway, so we'll just emit a
				// log to aid with debugging.
				log.Printf("[ERROR] Failed to persist state after interruption: %s", err)
			}
		} else {
			log.Printf("[DEBUG] State storage %T declined to persist a state snapshot", h.StateMgr)
		}
	}

}

func (h *StateHook) shouldPersist() bool {
	if m, ok := h.StateMgr.(IntermediateStateConditionalPersister); ok {
		return m.ShouldPersistIntermediateState(&h.intermediatePersist)
	}
	return DefaultIntermediateStatePersistRule(&h.intermediatePersist)
}

// DefaultIntermediateStatePersistRule is the default implementation of
// [IntermediateStateConditionalPersister.ShouldPersistIntermediateState] used
// when the selected state manager doesn't implement that interface.
//
// Implementers of that interface can optionally wrap a call to this function
// if they want to combine the default behavior with some logic of their own.
func DefaultIntermediateStatePersistRule(info *IntermediateStatePersistInfo) bool {
	return info.ForcePersist || time.Since(info.LastPersist) >= info.RequestedPersistInterval
}

// IntermediateStateConditionalPersister is an optional extension of
// [statemgr.Persister] that allows an implementation to tailor the rules for
// whether to create intermediate state snapshots when OpenTofu Core emits
// events reporting that the state might have changed.
//
// For state managers that don't implement this interface, [StateHook] uses
// a default set of rules that aim to be a good compromise between how long
// a state change can be active before it gets committed as a snapshot vs.
// how many intermediate snapshots will get created. That compromise is subject
// to change over time, but a state manager can implement this interface to
// exert full control over those rules.
type IntermediateStateConditionalPersister interface {
	// ShouldPersistIntermediateState will be called each time OpenTofu Core
	// emits an intermediate state event that is potentially eligible to be
	// persisted.
	//
	// The implemention should return true to signal that the state snapshot
	// most recently provided to the object's WriteState should be persisted,
	// or false if it should not be persisted. If this function returns true
	// then the receiver will see a subsequent call to
	// [statemgr.Persister.PersistState] to request persistence.
	//
	// The implementation must not modify anything reachable through the
	// arguments, and must not retain pointers to anything reachable through
	// them after the function returns. However, implementers can assume that
	// nothing will write to anything reachable through the arguments while
	// this function is active.
	ShouldPersistIntermediateState(info *IntermediateStatePersistInfo) bool
}