plugin-barman-cloud/internal/cnpgi/common/wal.go

/*
Copyright © contributors to CloudNativePG, established as
CloudNativePG a Series of LF Projects, LLC.

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.

SPDX-License-Identifier: Apache-2.0
*/

package common

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path"
	"sync"
	"time"

	barmanapi "github.com/cloudnative-pg/barman-cloud/pkg/api"
	"github.com/cloudnative-pg/barman-cloud/pkg/archiver"
	barmanCommand "github.com/cloudnative-pg/barman-cloud/pkg/command"
	barmanCredentials "github.com/cloudnative-pg/barman-cloud/pkg/credentials"
	barmanRestorer "github.com/cloudnative-pg/barman-cloud/pkg/restorer"
	cnpgv1 "github.com/cloudnative-pg/cloudnative-pg/api/v1"
	"github.com/cloudnative-pg/cloudnative-pg/pkg/utils"
	"github.com/cloudnative-pg/cnpg-i/pkg/wal"
	"github.com/cloudnative-pg/machinery/pkg/fileutils"
	walUtils "github.com/cloudnative-pg/machinery/pkg/fileutils/wals"
	"github.com/cloudnative-pg/machinery/pkg/log"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/util/retry"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	barmancloudv1 "github.com/cloudnative-pg/plugin-barman-cloud/api/v1"
	"github.com/cloudnative-pg/plugin-barman-cloud/internal/cnpgi/metadata"
	"github.com/cloudnative-pg/plugin-barman-cloud/internal/cnpgi/operator/config"
)

// SpoolManagementError is raised when a spool management
// error has been detected
type SpoolManagementError struct {
	walName string
	err     error
}

// Error implements the error interface
func (e *SpoolManagementError) Error() string {
	return fmt.Sprintf(
		"while testing the existence of the WAL file (%s) in the spool directory: %v",
		e.walName,
		e.err.Error(),
	)
}

// Unwrap implements the error interface
func (e *SpoolManagementError) Unwrap() error {
	return e.err
}

const (
	// walStatusUpdateThrottle is the minimum time between status updates for WAL archiving
	walStatusUpdateThrottle = 5 * time.Minute
)

// WALServiceImplementation is the implementation of the WAL Service
type WALServiceImplementation struct {
	wal.UnimplementedWALServer
	Client         client.Client
	InstanceName   string
	SpoolDirectory string
	PGDataPath     string
	PGWALPath      string
	// LastStatusUpdate tracks the last time we updated the status for each ObjectStore+ServerName
	// Key format: "namespace/objectStoreName/serverName"
	LastStatusUpdate *sync.Map
}

// GetCapabilities implements the WALService interface
func (w WALServiceImplementation) GetCapabilities(
	_ context.Context,
	_ *wal.WALCapabilitiesRequest,
) (*wal.WALCapabilitiesResult, error) {
	return &wal.WALCapabilitiesResult{
		Capabilities: []*wal.WALCapability{
			{
				Type: &wal.WALCapability_Rpc{
					Rpc: &wal.WALCapability_RPC{
						Type: wal.WALCapability_RPC_TYPE_ARCHIVE_WAL,
					},
				},
			},
			{
				Type: &wal.WALCapability_Rpc{
					Rpc: &wal.WALCapability_RPC{
						Type: wal.WALCapability_RPC_TYPE_RESTORE_WAL,
					},
				},
			},
		},
	}, nil
}

// shouldUpdateStatus checks if we should update the status based on the throttle.
// It returns true if walStatusUpdateThrottle minutes have passed since the last update, or if this is the first update.
func (w WALServiceImplementation) shouldUpdateStatus(objectStoreKey client.ObjectKey, serverName string) bool {
	if w.LastStatusUpdate == nil {
		return true
	}

	key := fmt.Sprintf("%s/%s", objectStoreKey.String(), serverName)
	lastUpdate, ok := w.LastStatusUpdate.Load(key)
	if !ok {
		return true
	}

	lastUpdateTime, ok := lastUpdate.(time.Time)
	if !ok {
		return true
	}

	return time.Since(lastUpdateTime) >= walStatusUpdateThrottle
}

// recordStatusUpdate records that we just updated the status for a given ObjectStore and server.
func (w WALServiceImplementation) recordStatusUpdate(objectStoreKey client.ObjectKey, serverName string) {
	if w.LastStatusUpdate == nil {
		return
	}

	key := fmt.Sprintf("%s/%s", objectStoreKey.String(), serverName)
	w.LastStatusUpdate.Store(key, time.Now())
}

// Archive implements the WALService interface
func (w WALServiceImplementation) Archive(
	ctx context.Context,
	request *wal.WALArchiveRequest,
) (*wal.WALArchiveResult, error) {
	baseWalName := path.Base(request.GetSourceFileName())

	contextLogger := log.FromContext(ctx)
	contextLogger.Debug("wal archive start", "walName", baseWalName)
	defer func() {
		contextLogger.Debug("wal archive end", "walName", baseWalName)
	}()

	// Step 1: parse the configuration and get the environment variables needed
	// for barman-cloud-wal-archive
	configuration, err := config.NewFromClusterJSON(request.ClusterDefinition)
	if err != nil {
		return nil, err
	}

	var objectStore barmancloudv1.ObjectStore
	if err := w.Client.Get(ctx, configuration.GetBarmanObjectKey(), &objectStore); err != nil {
		return nil, err
	}

	envArchive, err := barmanCredentials.EnvSetCloudCredentialsAndCertificates(
		ctx,
		w.Client,
		objectStore.Namespace,
		&objectStore.Spec.Configuration,
		os.Environ(),
		BuildCertificateFilePath(objectStore.Name),
	)
	if err != nil {
		if apierrors.IsForbidden(err) {
			contextLogger.Info(ErrMissingPermissions.Error())
			return nil, ErrMissingPermissions
		}
		return nil, err
	}

	emptyWalArchiveFile := path.Join(w.PGDataPath, metadata.CheckEmptyWalArchiveFile)
	arch, err := archiver.New(
		ctx,
		envArchive,
		w.SpoolDirectory,
		w.PGDataPath,
		emptyWalArchiveFile,
	)
	if err != nil {
		return nil, err
	}

	// Step 2: Check if the archive location is safe to perform archiving.
	checkEmptyWalArchive, err := resolveArchiveEmptyWalArchiveCheck(
		request.CheckEmptyWalArchive,
		configuration.Cluster,
		emptyWalArchiveFile,
	)
	if err != nil {
		return nil, err
	}

	if checkEmptyWalArchive {
		if err := CheckBackupDestination(
			ctx,
			&objectStore.Spec.Configuration,
			arch,
			configuration.ServerName,
		); err != nil {
			return nil, err
		}
	}

	// Step 3: check if this WAL file has not been already archived
	var isDeletedFromSpool bool
	isDeletedFromSpool, err = arch.DeleteFromSpool(baseWalName)
	if err != nil {
		return nil, &SpoolManagementError{
			walName: baseWalName,
			err:     err,
		}
	}
	if isDeletedFromSpool {
		contextLogger.Info("WAL file already archived, skipping",
			"walName", baseWalName)
		return nil, nil
	}

	// Step 4: gather the WAL files names to archive
	options, err := arch.BarmanCloudWalArchiveOptions(ctx, &objectStore.Spec.Configuration, configuration.ServerName)
	if err != nil {
		return nil, err
	}

	maxParallel := 1
	if objectStore.Spec.Configuration.Wal != nil && objectStore.Spec.Configuration.Wal.MaxParallel > 0 {
		maxParallel = objectStore.Spec.Configuration.Wal.MaxParallel
	}

	maxResults := maxParallel - 1
	walFilesList := walUtils.GatherReadyWALFiles(
		ctx,
		walUtils.GatherReadyWALFilesConfig{
			MaxResults: maxResults,
			SkipWALs:   []string{baseWalName},
			PgDataPath: w.PGDataPath,
		},
	)

	// Ensure the requested WAL file is always the first one being
	// archived
	walFilesList.Ready = append([]string{request.GetSourceFileName()}, walFilesList.Ready...)
	contextLogger.Debug("WAL files to archive", "walFilesListReady", walFilesList.Ready)

	result := arch.ArchiveList(ctx, walFilesList.ReadyItemsToSlice(), options)
	for _, archiverResult := range result {
		if archiverResult.Err != nil {
			return nil, archiverResult.Err
		}
	}

	// Update the last archived WAL time in the ObjectStore status
	// Only update if walStatusUpdateThrottle minutes have passed since the last update to avoid hitting the API server too often
	objectStoreKey := configuration.GetBarmanObjectKey()
	if w.shouldUpdateStatus(objectStoreKey, configuration.ServerName) {
		contextLogger.Debug("Updating last archived WAL time", "serverName", configuration.ServerName)
		if err := setLastArchivedWALTime(
			ctx,
			w.Client,
			objectStoreKey,
			configuration.ServerName,
			time.Now(),
		); err != nil {
			// Log the error but don't fail the archive operation
			contextLogger.Error(err, "Error updating last archived WAL time in ObjectStore status")
		} else {
			contextLogger.Debug("Successfully updated last archived WAL time")
			w.recordStatusUpdate(objectStoreKey, configuration.ServerName)
		}
	} else {
		contextLogger.Debug("Skipping status update due to throttle", "serverName", configuration.ServerName)
	}

	return &wal.WALArchiveResult{}, nil
}

// resolveArchiveEmptyWalArchiveCheck reports whether the WAL archive
// destination must be verified before archiving this segment.
//
// The operator owns the marker file's lifecycle, so when it sets the decision
// (non-nil) that value already accounts for the marker and is obeyed as-is. A
// nil value comes from an operator that predates this field, so we fall back to
// the previous logic: the Cluster annotation combined with the on-disk marker
// file.
func resolveArchiveEmptyWalArchiveCheck(
	operatorDecision *bool,
	cluster *cnpgv1.Cluster,
	markerFilePath string,
) (bool, error) {
	if operatorDecision != nil {
		return *operatorDecision, nil
	}

	markerFilePresent, err := fileutils.FileExists(markerFilePath)
	if err != nil {
		return false, fmt.Errorf("while checking for empty wal archive check file %q: %w", markerFilePath, err)
	}

	return utils.IsEmptyWalArchiveCheckEnabled(&cluster.ObjectMeta) && markerFilePresent, nil
}

// Restore implements the WALService interface
func (w WALServiceImplementation) Restore(
	ctx context.Context,
	request *wal.WALRestoreRequest,
) (*wal.WALRestoreResult, error) {
	contextLogger := log.FromContext(ctx)

	walName := request.GetSourceWalName()
	destinationPath := request.GetDestinationFileName()

	configuration, err := config.NewFromClusterJSON(request.ClusterDefinition)
	if err != nil {
		return nil, err
	}

	serverName, objectStoreKey := resolveRestoreObjectStore(configuration, w.InstanceName)

	var objectStore barmancloudv1.ObjectStore
	if err := w.Client.Get(ctx, objectStoreKey, &objectStore); err != nil {
		return nil, err
	}

	contextLogger.Info(
		"Restoring WAL file",
		"objectStore", objectStore.Name,
		"serverName", serverName,
		"walName", walName,
		"mode", request.GetMode())
	return &wal.WALRestoreResult{}, w.restoreFromBarmanObjectStore(
		ctx, configuration.Cluster, &objectStore, serverName, walName, destinationPath,
		request.GetMode() == wal.WALRestoreRequest_MODE_REWIND)
}

// resolveRestoreObjectStore selects the object store and server name to use when
// restoring a WAL file, based on the role this instance plays in the cluster.
func resolveRestoreObjectStore(
	configuration *config.PluginConfiguration,
	instanceName string,
) (serverName string, objectStoreKey types.NamespacedName) {
	switch {
	case configuration.Cluster.Status.CurrentPrimary == instanceName &&
		len(configuration.ReplicaSourceBarmanObjectName) > 0:
		// PostgreSQL never runs restore_command on a live primary; it runs only while
		// the instance is in recovery (a genuine standby, restoring from a backup, or
		// being rewound by pg_rewind). So a current primary that still has a replica
		// source configured can only be a designated primary that has not finished
		// promoting, and it must keep fetching WAL from the replica source.
		// Token-agnostic: covers both switchover and failover.
		return configuration.ReplicaSourceServerName, configuration.GetReplicaSourceBarmanObjectKey()

	case configuration.Cluster.Status.CurrentPrimary == "":
		// Recovery from object store, using recovery object store
		return configuration.RecoveryServerName, configuration.GetRecoveryBarmanObjectKey()

	default:
		// Using cluster object store
		return configuration.ServerName, configuration.GetBarmanObjectKey()
	}
}

func (w WALServiceImplementation) restoreFromBarmanObjectStore(
	ctx context.Context,
	cluster *cnpgv1.Cluster,
	objectStore *barmancloudv1.ObjectStore,
	serverName string,
	walName string,
	destinationPath string,
	rewindMode bool,
) error {
	contextLogger := log.FromContext(ctx)
	startTime := time.Now()

	barmanConfiguration := &objectStore.Spec.Configuration

	env := GetRestoreCABundleEnv(barmanConfiguration)
	credentialsEnv, err := barmanCredentials.EnvSetCloudCredentialsAndCertificates(
		ctx,
		w.Client,
		objectStore.Namespace,
		&objectStore.Spec.Configuration,
		os.Environ(),
		BuildCertificateFilePath(objectStore.Name),
	)
	if err != nil {
		return fmt.Errorf("while getting recover credentials: %w", err)
	}
	env = MergeEnv(env, credentialsEnv)

	options, err := barmanCommand.CloudWalRestoreOptions(ctx, barmanConfiguration, serverName)
	if err != nil {
		return fmt.Errorf("while getting barman-cloud-wal-restore options: %w", err)
	}

	// Create the restorer
	var walRestorer *barmanRestorer.WALRestorer
	if walRestorer, err = barmanRestorer.New(ctx, env, w.SpoolDirectory); err != nil {
		return fmt.Errorf("while creating the restorer: %w", err)
	}

	// A flag left over from a normal-recovery invocation that ran before this pod
	// was demoted must not survive into a pg_rewind restore: the flag machinery
	// does not apply while restoring on behalf of pg_rewind (see
	// shouldUseEndOfWALStreamFlag below), so it is never checked here, but left
	// untouched it would resurface and wrongly abort the first normal-recovery
	// invocation that runs once the rewind is done. This runs unconditionally,
	// before the spool short-circuit in Step 1, so a request for a WAL file that
	// happens to already be staged in the spool cannot skip the clear.
	if rewindMode {
		if err := clearEndOfWALStreamFlag(walRestorer); err != nil {
			return err
		}
	}

	// Step 1: check if this WAL file is not already in the spool
	var wasInSpool bool
	if wasInSpool, err = walRestorer.RestoreFromSpool(walName, destinationPath); err != nil {
		return fmt.Errorf("while restoring a file from the spool directory: %w", err)
	}
	if wasInSpool {
		contextLogger.Info("Restored WAL file from spool (parallel)",
			"walName", walName,
		)
		return nil
	}

	// Step 2: return error if the end-of-wal-stream flag is set.
	// We skip this step if the flag machinery does not apply to this invocation
	useEndOfWALStreamFlag := shouldUseEndOfWALStreamFlag(cluster, w.InstanceName, rewindMode)
	if useEndOfWALStreamFlag {
		if err := checkEndOfWALStreamFlag(walRestorer); err != nil {
			return err
		}
	}

	// Step 3: gather the WAL files names to restore. If the required file isn't a regular WAL, we download it directly.
	var walFilesList []string
	maxParallel := maxWALFilesPerInvocation(barmanConfiguration, rewindMode)
	if IsWALFile(walName) {
		// If this is a regular WAL file, we try to prefetch
		if walFilesList, err = gatherWALFilesToRestore(walName, maxParallel); err != nil {
			return fmt.Errorf("while generating the list of WAL files to restore: %w", err)
		}
	} else {
		// This is not a regular WAL file, we fetch it directly
		walFilesList = []string{walName}
	}

	// Step 4: download the WAL files into the required place
	downloadStartTime := time.Now()
	walStatus := walRestorer.RestoreList(ctx, walFilesList, destinationPath, options)

	// We return immediately if the first WAL has errors, because the first WAL
	// is the one that PostgreSQL has requested to restore.
	// The failure has already been logged in walRestorer.RestoreList method
	if walStatus[0].Err != nil {
		return classifyWALRestoreError(walStatus[0].WalName, walStatus[0].Err)
	}

	// We skip this step if the flag machinery does not apply to this invocation
	endOfWALStream := isEndOfWALStream(walStatus)
	if useEndOfWALStreamFlag && endOfWALStream {
		contextLogger.Info(
			"Set end-of-wal-stream flag as one of the WAL files to be prefetched was not found")

		err = walRestorer.SetEndOfWALStream()
		if err != nil {
			return err
		}
	}

	successfulWalRestore := 0
	for idx := range walStatus {
		if walStatus[idx].Err == nil {
			successfulWalRestore++
		}
	}

	contextLogger.Info("WAL restore command completed (parallel)",
		"walName", walName,
		"maxParallel", maxParallel,
		"successfulWalRestore", successfulWalRestore,
		"failedWalRestore", maxParallel-successfulWalRestore,
		"startTime", startTime,
		"downloadStartTime", downloadStartTime,
		"downloadTotalTime", time.Since(downloadStartTime),
		"totalTime", time.Since(startTime))

	return nil
}

// Status implements the WALService interface
func (w WALServiceImplementation) Status(
	_ context.Context,
	_ *wal.WALStatusRequest,
) (*wal.WALStatusResult, error) {
	// TODO implement me
	panic("implement me")
}

// SetFirstRequired implements the WALService interface
func (w WALServiceImplementation) SetFirstRequired(
	_ context.Context,
	_ *wal.SetFirstRequiredRequest,
) (*wal.SetFirstRequiredResult, error) {
	// TODO implement me
	panic("implement me")
}

// maxWALFilesPerInvocation returns how many WAL files a single restore
// invocation is allowed to fetch, the requested one included. Prefetching is
// disabled when restoring on behalf of pg_rewind, which walks the WAL
// backward: the following segments would never be requested, and past the end
// of the timeline they do not even exist
func maxWALFilesPerInvocation(barmanConfiguration *barmanapi.BarmanObjectStoreConfiguration, rewindMode bool) int {
	if rewindMode {
		return 1
	}

	if barmanConfiguration.Wal != nil && barmanConfiguration.Wal.MaxParallel > 1 {
		return barmanConfiguration.Wal.MaxParallel
	}

	return 1
}

// shouldUseEndOfWALStreamFlag returns true when the end-of-wal-stream flag
// machinery applies to the current invocation. The flag makes the following
// invocation fail, so that PostgreSQL stops polling the WAL archive and
// switches to streaming replication. It does not apply when no streaming
// connection is available, nor when restoring on behalf of pg_rewind:
// pg_rewind cannot fall back to streaming replication, and a stale flag would
// make it abort on a segment that is available in the archive
func shouldUseEndOfWALStreamFlag(cluster *cnpgv1.Cluster, podName string, rewindMode bool) bool {
	if rewindMode {
		return false
	}

	return isStreamingAvailable(cluster, podName)
}

// isStreamingAvailable checks if this pod can replicate via streaming connection.
func isStreamingAvailable(cluster *cnpgv1.Cluster, podName string) bool {
	if cluster == nil {
		return false
	}

	// Easy case take 1: we are helping PostgreSQL to create the first
	// instance of a Cluster. No streaming connection is possible.
	if cluster.Status.CurrentPrimary == "" {
		return false
	}

	// Easy case take 2: If this pod is a replica, the streaming is always
	// available
	if cluster.Status.CurrentPrimary != podName {
		return true
	}

	// Designated primary in a replica cluster: return true if the external cluster has streaming connection
	if cluster.IsReplica() {
		externalCluster, found := cluster.ExternalCluster(cluster.Spec.ReplicaCluster.Source)

		// This is a configuration error
		if !found {
			return false
		}

		return externalCluster.ConnectionParameters != nil
	}

	// Primary, we do not replicate from nobody
	return false
}

// gatherWALFilesToRestore files a list of possible WAL files to restore, always
// including as the first one the requested WAL file.
func gatherWALFilesToRestore(walName string, parallel int) (walList []string, err error) {
	var segment Segment

	segment, err = SegmentFromName(walName)
	if err != nil {
		// This seems an invalid segment name. It's not a problem
		// because PostgreSQL may request also other files such as
		// backup, history, etc.
		// Let's just avoid prefetching in this case
		return []string{walName}, nil
	}
	// NextSegments would accept postgresVersion and segmentSize,
	// but we do not have this info here, so we pass nil.
	segmentList := segment.NextSegments(parallel, nil, nil)
	walList = make([]string, len(segmentList))
	for idx := range segmentList {
		walList[idx] = segmentList[idx].Name()
	}

	return walList, err
}

// checkEndOfWALStreamFlag returns ErrEndOfWALStreamReached if the flag is set in the restorer.
func checkEndOfWALStreamFlag(walRestorer *barmanRestorer.WALRestorer) error {
	contain, err := walRestorer.IsEndOfWALStream()
	if err != nil {
		return err
	}

	if contain {
		err := walRestorer.ResetEndOfWalStream()
		if err != nil {
			return err
		}

		return ErrEndOfWALStreamReached
	}
	return nil
}

// clearEndOfWALStreamFlag removes the end-of-wal-stream flag, if present, without
// treating it as an error. It is used instead of checkEndOfWALStreamFlag when
// restoring on behalf of pg_rewind, which must not abort on a flag it did not
// set itself.
func clearEndOfWALStreamFlag(walRestorer *barmanRestorer.WALRestorer) error {
	contain, err := walRestorer.IsEndOfWALStream()
	if err != nil {
		return err
	}

	if contain {
		return walRestorer.ResetEndOfWalStream()
	}
	return nil
}

// isEndOfWALStream returns true if one of the downloads has returned
// a file-not-found error.
func isEndOfWALStream(results []barmanRestorer.Result) bool {
	for _, result := range results {
		if errors.Is(result.Err, barmanRestorer.ErrWALNotFound) {
			return true
		}
	}

	return false
}

// SetLastArchivedWALTime sets the last archived WAL time in the
// passed object store, for the passed server name.
func setLastArchivedWALTime(
	ctx context.Context,
	c client.Client,
	objectStoreKey client.ObjectKey,
	serverName string,
	lastArchivedWALTime time.Time,
) error {
	return retry.RetryOnConflict(retry.DefaultBackoff, func() error {
		var objectStore barmancloudv1.ObjectStore

		if err := c.Get(ctx, objectStoreKey, &objectStore); err != nil {
			return err
		}
		recoveryWindow := objectStore.Status.ServerRecoveryWindow[serverName]
		recoveryWindow.LastArchivedWALTime = ptr.To(metav1.NewTime(lastArchivedWALTime))

		if objectStore.Status.ServerRecoveryWindow == nil {
			objectStore.Status.ServerRecoveryWindow = make(map[string]barmancloudv1.RecoveryWindow)
		}
		objectStore.Status.ServerRecoveryWindow[serverName] = recoveryWindow

		return c.Status().Update(ctx, &objectStore)
	})
}