github.com/Schaudge/grailbase@v0.0.0-20240223061707-44c758a471c0/cloud/ec2util/iid.go (about)

     1  package ec2util
     2  
     3  import (
     4  	"fmt"
     5  	"net/http"
     6  	"time"
     7  
     8  	"github.com/aws/aws-sdk-go/aws"
     9  	"github.com/aws/aws-sdk-go/aws/ec2metadata"
    10  	"github.com/aws/aws-sdk-go/aws/session"
    11  	"github.com/Schaudge/grailbase/sync/once"
    12  )
    13  
    14  var (
    15  	iidOnce once.Task
    16  	iid     ec2metadata.EC2InstanceIdentityDocument
    17  )
    18  
    19  // GetInstanceIdentityDocument returns the EC2 Instance ID document (if the current
    20  // process is running within an EC2 instance) or an error.
    21  // Unlike the SDK's implementation, this will use longer timeouts and multiple retries
    22  // to improve the reliability of getting the Instance ID document.
    23  // The first result, whether success or failure, is cached for the lifetime of the process.
    24  func GetInstanceIdentityDocument(sess *session.Session) (doc ec2metadata.EC2InstanceIdentityDocument, err error) {
    25  	err = iidOnce.Do(func() (oerr error) {
    26  		// Use HTTP client with custom timeout and max retries to prevent the SDK from
    27  		// using an HTTP client with a small timeout and a small number of retries for
    28  		// the ec2metadata client
    29  		metaClient := ec2metadata.New(sess, &aws.Config{
    30  			HTTPClient: &http.Client{Timeout: 5 * time.Second},
    31  			MaxRetries: aws.Int(5),
    32  		})
    33  		for retries := 0; retries < 5; retries++ {
    34  			iid, oerr = metaClient.GetInstanceIdentityDocument()
    35  			if oerr == nil {
    36  				break
    37  			}
    38  		}
    39  		return
    40  	})
    41  	switch {
    42  	case err != nil:
    43  		err = fmt.Errorf("ec2util.GetInstanceIdentityDocument: %v", err)
    44  	case iid.InstanceID == "":
    45  		err = fmt.Errorf("ec2util.GetInstanceIdentityDocument: Unable to get EC2InstanceIdentityDocument")
    46  	default:
    47  		doc = iid
    48  	}
    49  	return
    50  }