postgres-operator/pkg/util/volumes/ebs.go

package volumes

import (
	"context"
	"fmt"
	"strings"

	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/service/ec2"
	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
	v1 "k8s.io/api/core/v1"

	"github.com/zalando/postgres-operator/pkg/util/constants"
	"github.com/zalando/postgres-operator/pkg/util/retryutil"
)

// EBSVolumeResizer implements volume resizing interface for AWS EBS volumes.
type EBSVolumeResizer struct {
	connection *ec2.Client
	AWSRegion  string
}

// ConnectToProvider connects to AWS.
func (r *EBSVolumeResizer) ConnectToProvider() error {
	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion(r.AWSRegion))
	if err != nil {
		return fmt.Errorf("could not establish AWS session: %v", err)
	}
	r.connection = ec2.NewFromConfig(cfg)
	return nil
}

// IsConnectedToProvider checks if AWS connection is established.
func (r *EBSVolumeResizer) IsConnectedToProvider() bool {
	return r.connection != nil
}

// VolumeBelongsToProvider checks if the given persistent volume is backed by EBS.
func (r *EBSVolumeResizer) VolumeBelongsToProvider(pv *v1.PersistentVolume) bool {
	return (pv.Spec.AWSElasticBlockStore != nil && pv.Annotations[constants.VolumeStorateProvisionerAnnotation] == constants.EBSProvisioner) ||
		(pv.Spec.CSI != nil && pv.Spec.CSI.Driver == constants.EBSDriver)
}

// ExtractVolumeID extracts volumeID from "aws://eu-central-1a/vol-075ddfc4a127d0bd4"
// or return only the vol-075ddfc4a127d0bd4 when it doesn't have "aws://"
func (r *EBSVolumeResizer) ExtractVolumeID(volumeID string) (string, error) {
	if (strings.HasPrefix(volumeID, "vol-")) && !(strings.HasPrefix(volumeID, "aws://")) {
		return volumeID, nil
	}
	idx := strings.LastIndex(volumeID, constants.EBSVolumeIDStart) + 1
	if idx == 0 {
		return "", fmt.Errorf("malformed EBS volume id %q", volumeID)
	}
	return volumeID[idx:], nil
}

// GetProviderVolumeID converts aws://eu-central-1b/vol-00f93d4827217c629 to vol-00f93d4827217c629 for EBS volumes
func (r *EBSVolumeResizer) GetProviderVolumeID(pv *v1.PersistentVolume) (string, error) {
	var volumeID string = ""
	if pv.Spec.CSI != nil {
		volumeID = pv.Spec.CSI.VolumeHandle
	} else if pv.Spec.AWSElasticBlockStore != nil {
		volumeID = pv.Spec.AWSElasticBlockStore.VolumeID
	}
	if volumeID == "" {
		return "", fmt.Errorf("got empty volume id for volume %v", pv)
	}

	return r.ExtractVolumeID(volumeID)
}

// DescribeVolumes ...
func (r *EBSVolumeResizer) DescribeVolumes(volumeIds []string) ([]VolumeProperties, error) {
	if !r.IsConnectedToProvider() {
		err := r.ConnectToProvider()
		if err != nil {
			return nil, err
		}
	}

	volumeOutput, err := r.connection.DescribeVolumes(context.TODO(), &ec2.DescribeVolumesInput{VolumeIds: volumeIds})
	if err != nil {
		return nil, err
	}

	p := []VolumeProperties{}
	if nil == volumeOutput.Volumes {
		return p, nil
	}

	for _, v := range volumeOutput.Volumes {
		tags := make(map[string]string)
		for _, tag := range v.Tags {
			if tag.Key != nil && tag.Value != nil {
				tags[*tag.Key] = *tag.Value
			}
		}

		switch v.VolumeType {
		case "gp3":
			p = append(p, VolumeProperties{VolumeID: *v.VolumeId, Size: int64(*v.Size), VolumeType: string(v.VolumeType), Iops: int64(*v.Iops), Throughput: int64(*v.Throughput), Tags: tags})
		case "gp2":
			p = append(p, VolumeProperties{VolumeID: *v.VolumeId, Size: int64(*v.Size), VolumeType: string(v.VolumeType), Tags: tags})
		default:
			return nil, fmt.Errorf("discovered unexpected volume type %s %s", *v.VolumeId, v.VolumeType)
		}
	}

	return p, nil
}

// ResizeVolume actually calls AWS API to resize the EBS volume if necessary.
func (r *EBSVolumeResizer) ResizeVolume(volumeID string, newSize int64) error {
	/* first check if the volume is already of a requested size */
	volumeOutput, err := r.connection.DescribeVolumes(context.TODO(), &ec2.DescribeVolumesInput{VolumeIds: []string{volumeID}})
	if err != nil {
		return fmt.Errorf("could not get information about the volume: %v", err)
	}
	vol := volumeOutput.Volumes[0]
	if *vol.VolumeId != volumeID {
		return fmt.Errorf("describe volume %q returned information about a non-matching volume %q", volumeID, *vol.VolumeId)
	}
	sizeInt32 := int32(newSize)
	if *vol.Size == sizeInt32 {
		// nothing to do
		return nil
	}
	input := ec2.ModifyVolumeInput{Size: &sizeInt32, VolumeId: &volumeID}
	output, err := r.connection.ModifyVolume(context.TODO(), &input)
	if err != nil {
		return fmt.Errorf("could not modify persistent volume: %v", err)
	}

	state := output.VolumeModification.ModificationState
	if state == constants.EBSVolumeStateFailed {
		return fmt.Errorf("could not modify persistent volume %q: modification state failed", volumeID)
	}
	if state == "" {
		return fmt.Errorf("received empty modification status")
	}
	if state == constants.EBSVolumeStateOptimizing || state == constants.EBSVolumeStateCompleted {
		return nil
	}
	// wait until the volume reaches the "optimizing" or "completed" state
	in := ec2.DescribeVolumesModificationsInput{VolumeIds: []string{volumeID}}
	return retryutil.Retry(constants.EBSVolumeResizeWaitInterval, constants.EBSVolumeResizeWaitTimeout,
		func() (bool, error) {
			out, err := r.connection.DescribeVolumesModifications(context.TODO(), &in)
			if err != nil {
				return false, fmt.Errorf("could not describe volume modification: %v", err)
			}
			if len(out.VolumesModifications) != 1 {
				return false, fmt.Errorf("describe volume modification didn't return one record for volume %q", volumeID)
			}
			if *out.VolumesModifications[0].VolumeId != volumeID {
				return false, fmt.Errorf("non-matching volume id when describing modifications: %q is different from %q",
					*out.VolumesModifications[0].VolumeId, volumeID)
			}
			return out.VolumesModifications[0].ModificationState != constants.EBSVolumeStateModifying, nil
		})
}

// ModifyVolume Modify EBS volume
func (r *EBSVolumeResizer) ModifyVolume(volumeID string, newType *string, newSize *int64, iops *int64, throughput *int64) error {
	/* first check if the volume is already of a requested size */
	var sizeInt32 *int32
	var iopsInt32 *int32
	var throughputInt32 *int32
	if newSize != nil {
		s := int32(*newSize)
		sizeInt32 = &s
	}
	if iops != nil {
		i := int32(*iops)
		iopsInt32 = &i
	}
	if throughput != nil {
		t := int32(*throughput)
		throughputInt32 = &t
	}
	input := ec2.ModifyVolumeInput{Size: sizeInt32, VolumeId: &volumeID, VolumeType: types.VolumeType(*newType), Iops: iopsInt32, Throughput: throughputInt32}
	output, err := r.connection.ModifyVolume(context.TODO(), &input)
	if err != nil {
		return fmt.Errorf("could not modify persistent volume: %v", err)
	}

	state := output.VolumeModification.ModificationState
	if state == constants.EBSVolumeStateFailed {
		return fmt.Errorf("could not modify persistent volume %q: modification state failed", volumeID)
	}
	if state == "" {
		return fmt.Errorf("received empty modification status")
	}
	if state == constants.EBSVolumeStateOptimizing || state == constants.EBSVolumeStateCompleted {
		return nil
	}
	// wait until the volume reaches the "optimizing" or "completed" state
	in := ec2.DescribeVolumesModificationsInput{VolumeIds: []string{volumeID}}
	return retryutil.Retry(constants.EBSVolumeResizeWaitInterval, constants.EBSVolumeResizeWaitTimeout,
		func() (bool, error) {
			out, err := r.connection.DescribeVolumesModifications(context.TODO(), &in)
			if err != nil {
				return false, fmt.Errorf("could not describe volume modification: %v", err)
			}
			if len(out.VolumesModifications) != 1 {
				return false, fmt.Errorf("describe volume modification didn't return one record for volume %q", volumeID)
			}
			if *out.VolumesModifications[0].VolumeId != volumeID {
				return false, fmt.Errorf("non-matching volume id when describing modifications: %q is different from %q",
					*out.VolumesModifications[0].VolumeId, volumeID)
			}
			return out.VolumesModifications[0].ModificationState != constants.EBSVolumeStateModifying, nil
		})
}

// TagVolumes tags the given EBS volumes with the provided tags.
// Callers are responsible for filtering out volumes that already have the desired tags.
func (r *EBSVolumeResizer) TagVolumes(volumeIds []string, tags map[string]string) error {
	if !r.IsConnectedToProvider() {
		err := r.ConnectToProvider()
		if err != nil {
			return err
		}
	}

	if len(volumeIds) == 0 {
		return nil
	}

	ec2Tags := make([]types.Tag, 0, len(tags))
	for key, value := range tags {
		ec2Tags = append(ec2Tags, types.Tag{Key: &key, Value: &value})
	}

	input := &ec2.CreateTagsInput{
		Resources: volumeIds,
		Tags:      ec2Tags,
	}

	_, err := r.connection.CreateTags(context.TODO(), input)
	if err != nil {
		return fmt.Errorf("could not tag EBS volumes: %v", err)
	}
	return nil
}

// DisconnectFromProvider closes connection to the EC2 instance
func (r *EBSVolumeResizer) DisconnectFromProvider() error {
	r.connection = nil
	return nil
}