From bc0e76c8a3378a677ab268b8afb608d43c9a8f45 Mon Sep 17 00:00:00 2001
From: George Brighton <gebn@users.noreply.github.com>
Date: Fri, 23 Jul 2021 08:42:03 +0100
Subject: [PATCH] Add AZ ID label to discovered EC2 targets (#8896)

* Add AZ ID to EC2 SD

Signed-off-by: George Brighton <george@gebn.co.uk>
---
 discovery/aws/ec2.go                | 48 ++++++++++++++++++++++++++---
 docs/configuration/configuration.md |  1 +
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/discovery/aws/ec2.go b/discovery/aws/ec2.go
index c306c920b..d2ad9736a 100644
--- a/discovery/aws/ec2.go
+++ b/discovery/aws/ec2.go
@@ -28,6 +28,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
 	"github.com/pkg/errors"
 	"github.com/prometheus/common/config"
 	"github.com/prometheus/common/model"
@@ -42,6 +43,7 @@ const (
 	ec2Label                  = model.MetaLabelPrefix + "ec2_"
 	ec2LabelAMI               = ec2Label + "ami"
 	ec2LabelAZ                = ec2Label + "availability_zone"
+	ec2LabelAZID              = ec2Label + "availability_zone_id"
 	ec2LabelArch              = ec2Label + "architecture"
 	ec2LabelIPv6Addresses     = ec2Label + "ipv6_addresses"
 	ec2LabelInstanceID        = ec2Label + "instance_id"
@@ -132,8 +134,14 @@ func (c *EC2SDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
 // the Discoverer interface.
 type EC2Discovery struct {
 	*refresh.Discovery
-	cfg *EC2SDConfig
-	ec2 *ec2.EC2
+	logger log.Logger
+	cfg    *EC2SDConfig
+	ec2    *ec2.EC2
+
+	// azToAZID maps this account's availability zones to their underlying AZ
+	// ID, e.g. eu-west-2a -> euw2-az2. Refreshes are performed sequentially, so
+	// no locking is required.
+	azToAZID map[string]string
 }
 
 // NewEC2Discovery returns a new EC2Discovery which periodically refreshes its targets.
@@ -142,7 +150,8 @@ func NewEC2Discovery(conf *EC2SDConfig, logger log.Logger) *EC2Discovery {
 		logger = log.NewNopLogger()
 	}
 	d := &EC2Discovery{
-		cfg: conf,
+		logger: logger,
+		cfg:    conf,
 	}
 	d.Discovery = refresh.NewDiscovery(
 		logger,
@@ -153,7 +162,7 @@ func NewEC2Discovery(conf *EC2SDConfig, logger log.Logger) *EC2Discovery {
 	return d
 }
 
-func (d *EC2Discovery) ec2Client() (*ec2.EC2, error) {
+func (d *EC2Discovery) ec2Client(ctx context.Context) (*ec2.EC2, error) {
 	if d.ec2 != nil {
 		return d.ec2, nil
 	}
@@ -185,8 +194,28 @@ func (d *EC2Discovery) ec2Client() (*ec2.EC2, error) {
 	return d.ec2, nil
 }
 
+func (d *EC2Discovery) azID(ctx context.Context, az string) (string, error) {
+	if azID, ok := d.azToAZID[az]; ok {
+		return azID, nil
+	}
+
+	azs, err := d.ec2.DescribeAvailabilityZonesWithContext(ctx, &ec2.DescribeAvailabilityZonesInput{})
+	if err != nil {
+		return "", errors.Wrap(err, "could not describe availability zones")
+	}
+	d.azToAZID = make(map[string]string, len(azs.AvailabilityZones))
+	for _, az := range azs.AvailabilityZones {
+		d.azToAZID[*az.ZoneName] = *az.ZoneId
+	}
+
+	if azID, ok := d.azToAZID[az]; ok {
+		return azID, nil
+	}
+	return "", fmt.Errorf("no availability zone ID mapping found for %s", az)
+}
+
 func (d *EC2Discovery) refresh(ctx context.Context) ([]*targetgroup.Group, error) {
-	ec2Client, err := d.ec2Client()
+	ec2Client, err := d.ec2Client(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -211,6 +240,14 @@ func (d *EC2Discovery) refresh(ctx context.Context) ([]*targetgroup.Group, error
 				if inst.PrivateIpAddress == nil {
 					continue
 				}
+				azID, err := d.azID(ctx, *inst.Placement.AvailabilityZone)
+				if err != nil {
+					level.Warn(d.logger).Log(
+						"msg", "Unable to determine availability zone ID",
+						"az", *inst.Placement.AvailabilityZone,
+						"err", err)
+				}
+
 				labels := model.LabelSet{
 					ec2LabelInstanceID: model.LabelValue(*inst.InstanceId),
 				}
@@ -237,6 +274,7 @@ func (d *EC2Discovery) refresh(ctx context.Context) ([]*targetgroup.Group, error
 
 				labels[ec2LabelAMI] = model.LabelValue(*inst.ImageId)
 				labels[ec2LabelAZ] = model.LabelValue(*inst.Placement.AvailabilityZone)
+				labels[ec2LabelAZID] = model.LabelValue(azID)
 				labels[ec2LabelInstanceState] = model.LabelValue(*inst.State.Name)
 				labels[ec2LabelInstanceType] = model.LabelValue(*inst.InstanceType)
 
diff --git a/docs/configuration/configuration.md b/docs/configuration/configuration.md
index 23e226d2c..8276ebbea 100644
--- a/docs/configuration/configuration.md
+++ b/docs/configuration/configuration.md
@@ -893,6 +893,7 @@ The following meta labels are available on targets during [relabeling](#relabel_
 * `__meta_ec2_ami`: the EC2 Amazon Machine Image
 * `__meta_ec2_architecture`: the architecture of the instance
 * `__meta_ec2_availability_zone`: the availability zone in which the instance is running
+* `__meta_ec2_availability_zone_id`: the [availability zone ID](https://docs.aws.amazon.com/ram/latest/userguide/working-with-az-ids.html) in which the instance is running
 * `__meta_ec2_instance_id`: the EC2 instance ID
 * `__meta_ec2_instance_lifecycle`: the lifecycle of the EC2 instance, set only for 'spot' or 'scheduled' instances, absent otherwise
 * `__meta_ec2_instance_state`: the state of the EC2 instance