src/pkg/gce/gce.go - cos/tools - Git at Google

 // Copyright 2018 Google 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.

 // Package gce contains high-level functionality for manipulating GCE resources.
 package gce

 import (
 	"context"
 	"errors"
 	"fmt"
 	"net/http"
 	"regexp"
 	"sort"
 	"strconv"
 	"strings"
 	"time"

 	"cos.googlesource.com/cos/tools.git/src/pkg/config"

 	compute "google.golang.org/api/compute/v1"
 	"google.golang.org/api/googleapi"
 )

 const (
 	defaultOperationTimeout = time.Duration(600) * time.Second
 	defaultRetryInterval    = time.Duration(5) * time.Second
 	gcsURLPrefix            = "https://storage.googleapis.com"
 )

 type timePkg struct {
 	Now   func() time.Time
 	Sleep func(time.Duration)
 }

 var (
 	// ErrTimeout indicates that an operation timed out.
 	ErrTimeout = errors.New("operation timed out")

 	// ErrImageNotFound indicates that a GCE image could not be found
 	ErrImageNotFound = errors.New("image not found")

 	realTime = &timePkg{time.Now, time.Sleep}

 	// This should match <prefix>-<channel>-<milestone>-<buildnumber>.
 	// This is the format of images in cos-cloud.
 	// Example: cos-dev-72-11172-0-0
 	imageNameRegex = regexp.MustCompile("[a-z0-9-]+-[a-z]+-([0-9]+)-([0-9]+-[0-9]+-[0-9]+)")
 )

 // buildDeprecationStatus constructs a *compute.DeprecationStatus struct used in a Deprecate GCE API
 // call. It fills in the structure with the "DEPRECATED" state, the given replacement, and the given
 // delete time, if provided.
 func buildDeprecationStatus(replacement string, deleteTime time.Time) *compute.DeprecationStatus {
 	status := &compute.DeprecationStatus{State: "DEPRECATED", Replacement: replacement}
 	if !deleteTime.IsZero() {
 		status.Deleted = deleteTime.Format(time.RFC3339)
 	}
 	return status
 }

 func waitForOp(svc *compute.Service, project string, op *compute.Operation, deadline time.Time, t *timePkg) error {
 	if op.Error != nil {
 		return fmt.Errorf("error with operation. name: %s error: %v", op.Name, op.Error)
 	}
 	if op.Status == "DONE" {
 		return nil
 	}
 	for {
 		t.Sleep(defaultRetryInterval)
 		op, err := svc.GlobalOperations.Get(project, op.Name).Do()
 		if err != nil {
 			return err
 		}
 		if op.Error != nil {
 			return fmt.Errorf("error with operation. name: %s error: %v", op.Name, op.Error)
 		}
 		if op.Status == "DONE" {
 			return nil
 		}
 		if t.Now().After(deadline) {
 			return ErrTimeout
 		}
 	}
 }

 func waitForOps(svc *compute.Service, project string, ops []*compute.Operation, t *timePkg) error {
 	deadline := t.Now().Add(defaultOperationTimeout)
 	for _, op := range ops {
 		if err := waitForOp(svc, project, op, deadline, t); err != nil {
 			return err
 		}
 	}
 	return nil
 }

 func deprecateInFamily(ctx context.Context, svc *compute.Service, newImage *config.Image, ttl int, t *timePkg) error {
 	if newImage.Family == "" {
 		return fmt.Errorf("input image does not have a family for deprecateInFamily. image: %v", newImage)
 	}
 	filter := fmt.Sprintf("(family = %s) (name != %s)", newImage.Family, newImage.Name)
 	images := []*compute.Image{}
 	err := svc.Images.List(newImage.Project).Filter(filter).Pages(ctx, func(imageList *compute.ImageList) error {
 		images = append(images, imageList.Items...)
 		return nil
 	})
 	if err != nil {
 		return err
 	}
 	ops := []*compute.Operation{}
 	for _, image := range images {
 		if image.Deprecated != nil {
 			continue
 		}
 		deleteTime := time.Time{}
 		if ttl > 0 {
 			deleteTime = t.Now().Add(time.Duration(ttl) * time.Second)
 		}
 		status := buildDeprecationStatus(newImage.URL(), deleteTime)
 		op, err := svc.Images.Deprecate(newImage.Project, image.Name, status).Do()
 		if err != nil {
 			return err
 		}
 		ops = append(ops, op)
 	}
 	return waitForOps(svc, newImage.Project, ops, t)
 }

 // DeprecateInFamily deprecates all of the old images in an image family.
 // Allows for assigning TTLs (in seconds) to deprecated images.
 func DeprecateInFamily(ctx context.Context, svc *compute.Service, newImage *config.Image, ttl int) error {
 	return deprecateInFamily(ctx, svc, newImage, ttl, realTime)
 }

 // ImageExists checks to see if the given image exists in the given project.
 func ImageExists(svc *compute.Service, project, name string) (bool, error) {
 	if _, err := svc.Images.Get(project, name).Do(); err != nil {
 		if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
 			return false, nil
 		}
 		return false, err
 	}
 	return true, nil
 }

 // CreateImage creates an image with imageName with the source-url from gcs storage
 func CreateImage(svc *compute.Service, sourceURL, imageName, imageProject string) error {
 	gcsImageURL := fmt.Sprintf("%s/%s", gcsURLPrefix, sourceURL)
 	image := &compute.Image{
 		Name: imageName,
 		RawDisk: &compute.ImageRawDisk{
 			Source: gcsImageURL,
 		},
 	}
 	createImageOp, err := svc.Images.Insert(imageProject, image).Do()
 	if err != nil {
 		return err
 	}
 	deadline := time.Now().Add(defaultOperationTimeout)
 	return waitForOp(svc, imageProject, createImageOp, deadline, realTime)
 }

 type decodedImageName struct {
 	name        string
 	milestone   int
 	buildNumber string
 }

 // newDecodedImageName decodes an image name from cos-cloud and returns
 // image information encoded in that image name.
 func newDecodedImageName(name string) (*decodedImageName, error) {
 	match := imageNameRegex.FindStringSubmatch(name)
 	if match == nil {
 		return nil, fmt.Errorf("could not parse name %s", name)
 	}
 	milestone, err := strconv.Atoi(match[1])
 	if err != nil {
 		return nil, fmt.Errorf("could not convert %s to a milestone: %s", match[1], err)
 	}
 	return &decodedImageName{name, milestone, match[2]}, nil
 }

 func imageCompare(first, second *decodedImageName) bool {
 	if first.milestone != second.milestone {
 		return first.milestone < second.milestone
 	}
 	for i := 0; i < 3; i++ {
 		// Because of how decodedImageNames are created (see newDecodedImageName),
 		// these atoi operations are guaranteed to work.
 		firstNum, _ := strconv.Atoi(strings.Split(first.buildNumber, "-")[i])
 		secondNum, _ := strconv.Atoi(strings.Split(second.buildNumber, "-")[i])
 		if firstNum != secondNum {
 			return firstNum < secondNum
 		}
 	}
 	return false
 }

 // ResolveMilestone gets the name of the latest COS image on the given milestone.
 // This resolution is done by looking at the image names in cos-cloud.
 func ResolveMilestone(ctx context.Context, svc *compute.Service, milestone int) (string, error) {
 	var images []*compute.Image
 	err := svc.Images.List("cos-cloud").Pages(ctx, func(imageList *compute.ImageList) error {
 		images = append(images, imageList.Items...)
 		return nil
 	})
 	if err != nil {
 		return "", err
 	}
 	var inMilestone []*decodedImageName
 	for _, image := range images {
 		decoded, err := newDecodedImageName(image.Name)
 		if err != nil {
 			continue
 		}
 		if decoded.milestone == milestone {
 			inMilestone = append(inMilestone, decoded)
 		}
 	}
 	if len(inMilestone) == 0 {
 		return "", ErrImageNotFound
 	}
 	sort.Slice(inMilestone, func(i, j int) bool {
 		return imageCompare(inMilestone[i], inMilestone[j])
 	})
 	return inMilestone[len(inMilestone)-1].name, nil
 }
	// Copyright 2018 Google 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.

	// Package gce contains high-level functionality for manipulating GCE resources.
	package gce

	import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"cos.googlesource.com/cos/tools.git/src/pkg/config"

	compute "google.golang.org/api/compute/v1"
	"google.golang.org/api/googleapi"
	)

	const (
	defaultOperationTimeout = time.Duration(600) * time.Second
	defaultRetryInterval = time.Duration(5) * time.Second
	gcsURLPrefix = "https://storage.googleapis.com"
	)

	type timePkg struct {
	Now func() time.Time
	Sleep func(time.Duration)
	}

	var (
	// ErrTimeout indicates that an operation timed out.
	ErrTimeout = errors.New("operation timed out")

	// ErrImageNotFound indicates that a GCE image could not be found
	ErrImageNotFound = errors.New("image not found")

	realTime = &timePkg{time.Now, time.Sleep}

	// This should match <prefix>-<channel>-<milestone>-<buildnumber>.
	// This is the format of images in cos-cloud.
	// Example: cos-dev-72-11172-0-0
	imageNameRegex = regexp.MustCompile("[a-z0-9-]+-[a-z]+-([0-9]+)-([0-9]+-[0-9]+-[0-9]+)")
	)

	// buildDeprecationStatus constructs a *compute.DeprecationStatus struct used in a Deprecate GCE API
	// call. It fills in the structure with the "DEPRECATED" state, the given replacement, and the given
	// delete time, if provided.
	func buildDeprecationStatus(replacement string, deleteTime time.Time) *compute.DeprecationStatus {
	status := &compute.DeprecationStatus{State: "DEPRECATED", Replacement: replacement}
	if !deleteTime.IsZero() {
	status.Deleted = deleteTime.Format(time.RFC3339)
	}
	return status
	}

	func waitForOp(svc compute.Service, project string, op compute.Operation, deadline time.Time, t *timePkg) error {
	if op.Error != nil {
	return fmt.Errorf("error with operation. name: %s error: %v", op.Name, op.Error)
	}
	if op.Status == "DONE" {
	return nil
	}
	for {
	t.Sleep(defaultRetryInterval)
	op, err := svc.GlobalOperations.Get(project, op.Name).Do()
	if err != nil {
	return err
	}
	if op.Error != nil {
	return fmt.Errorf("error with operation. name: %s error: %v", op.Name, op.Error)
	}
	if op.Status == "DONE" {
	return nil
	}
	if t.Now().After(deadline) {
	return ErrTimeout
	}
	}
	}

	func waitForOps(svc compute.Service, project string, ops []compute.Operation, t *timePkg) error {
	deadline := t.Now().Add(defaultOperationTimeout)
	for _, op := range ops {
	if err := waitForOp(svc, project, op, deadline, t); err != nil {
	return err
	}
	}
	return nil
	}

	func deprecateInFamily(ctx context.Context, svc compute.Service, newImage config.Image, ttl int, t *timePkg) error {
	if newImage.Family == "" {
	return fmt.Errorf("input image does not have a family for deprecateInFamily. image: %v", newImage)
	}
	filter := fmt.Sprintf("(family = %s) (name != %s)", newImage.Family, newImage.Name)
	images := []*compute.Image{}
	err := svc.Images.List(newImage.Project).Filter(filter).Pages(ctx, func(imageList *compute.ImageList) error {
	images = append(images, imageList.Items...)
	return nil
	})
	if err != nil {
	return err
	}
	ops := []*compute.Operation{}
	for _, image := range images {
	if image.Deprecated != nil {
	continue
	}
	deleteTime := time.Time{}
	if ttl > 0 {
	deleteTime = t.Now().Add(time.Duration(ttl) * time.Second)
	}
	status := buildDeprecationStatus(newImage.URL(), deleteTime)
	op, err := svc.Images.Deprecate(newImage.Project, image.Name, status).Do()
	if err != nil {
	return err
	}
	ops = append(ops, op)
	}
	return waitForOps(svc, newImage.Project, ops, t)
	}

	// DeprecateInFamily deprecates all of the old images in an image family.
	// Allows for assigning TTLs (in seconds) to deprecated images.
	func DeprecateInFamily(ctx context.Context, svc compute.Service, newImage config.Image, ttl int) error {
	return deprecateInFamily(ctx, svc, newImage, ttl, realTime)
	}

	// ImageExists checks to see if the given image exists in the given project.
	func ImageExists(svc *compute.Service, project, name string) (bool, error) {
	if _, err := svc.Images.Get(project, name).Do(); err != nil {
	if e, ok := err.(*googleapi.Error); ok && e.Code == http.StatusNotFound {
	return false, nil
	}
	return false, err
	}
	return true, nil
	}

	// CreateImage creates an image with imageName with the source-url from gcs storage
	func CreateImage(svc *compute.Service, sourceURL, imageName, imageProject string) error {
	gcsImageURL := fmt.Sprintf("%s/%s", gcsURLPrefix, sourceURL)
	image := &compute.Image{
	Name: imageName,
	RawDisk: &compute.ImageRawDisk{
	Source: gcsImageURL,
	},
	}
	createImageOp, err := svc.Images.Insert(imageProject, image).Do()
	if err != nil {
	return err
	}
	deadline := time.Now().Add(defaultOperationTimeout)
	return waitForOp(svc, imageProject, createImageOp, deadline, realTime)
	}

	type decodedImageName struct {
	name string
	milestone int
	buildNumber string
	}

	// newDecodedImageName decodes an image name from cos-cloud and returns
	// image information encoded in that image name.
	func newDecodedImageName(name string) (*decodedImageName, error) {
	match := imageNameRegex.FindStringSubmatch(name)
	if match == nil {
	return nil, fmt.Errorf("could not parse name %s", name)
	}
	milestone, err := strconv.Atoi(match[1])
	if err != nil {
	return nil, fmt.Errorf("could not convert %s to a milestone: %s", match[1], err)
	}
	return &decodedImageName{name, milestone, match[2]}, nil
	}

	func imageCompare(first, second *decodedImageName) bool {
	if first.milestone != second.milestone {
	return first.milestone < second.milestone
	}
	for i := 0; i < 3; i++ {
	// Because of how decodedImageNames are created (see newDecodedImageName),
	// these atoi operations are guaranteed to work.
	firstNum, _ := strconv.Atoi(strings.Split(first.buildNumber, "-")[i])
	secondNum, _ := strconv.Atoi(strings.Split(second.buildNumber, "-")[i])
	if firstNum != secondNum {
	return firstNum < secondNum
	}
	}
	return false
	}

	// ResolveMilestone gets the name of the latest COS image on the given milestone.
	// This resolution is done by looking at the image names in cos-cloud.
	func ResolveMilestone(ctx context.Context, svc *compute.Service, milestone int) (string, error) {
	var images []*compute.Image
	err := svc.Images.List("cos-cloud").Pages(ctx, func(imageList *compute.ImageList) error {
	images = append(images, imageList.Items...)
	return nil
	})
	if err != nil {
	return "", err
	}
	var inMilestone []*decodedImageName
	for _, image := range images {
	decoded, err := newDecodedImageName(image.Name)
	if err != nil {
	continue
	}
	if decoded.milestone == milestone {
	inMilestone = append(inMilestone, decoded)
	}
	}
	if len(inMilestone) == 0 {
	return "", ErrImageNotFound
	}
	sort.Slice(inMilestone, func(i, j int) bool {
	return imageCompare(inMilestone[i], inMilestone[j])
	})
	return inMilestone[len(inMilestone)-1].name, nil
	}