/
aws.go
66 lines (58 loc) · 1.79 KB
/
aws.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
package kafka_queue
// based on https://github.com/segmentio/kafka-go/blob/e0af1cfbb8dd463571748e350262e2c81754bb73/example_groupbalancer_test.go#L37 and https://github.com/segmentio/kafka-go/issues/415
import (
"encoding/json"
"net/http"
"os"
"time"
)
var availabilityZoneMap = map[string]string{
"us-east-2a": "use2-az1",
"us-east-2b": "use2-az2",
"us-east-2c": "use2-az3",
}
// findRack is the basic rack resolver strategy for use in AWS. It supports
// - ECS with the task metadata endpoint enabled (returns the container
// instance's availability zone)
// - Linux EC2 (returns the instance's availability zone)
func findRack() string {
switch whereAmI() {
case "ecs":
return ecsAvailabilityZone()
}
return ""
}
const ecsContainerMetadataURI = "ECS_CONTAINER_METADATA_URI_V4"
// whereAmI determines which strategy the rack resolver should use.
func whereAmI() string {
// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint.html
if os.Getenv(ecsContainerMetadataURI) != "" {
return "ecs"
}
return "somewhere"
}
// ecsAvailabilityZone queries the task endpoint for the metadata URI that ECS
// injects into the ECS_CONTAINER_METADATA_URI variable in order to retrieve
// the availability zone where the task is running.
func ecsAvailabilityZone() string {
client := http.Client{
Timeout: time.Second,
Transport: &http.Transport{
DisableCompression: true,
DisableKeepAlives: true,
},
}
r, err := client.Get(os.Getenv(ecsContainerMetadataURI) + "/task")
if err != nil {
return ""
}
defer r.Body.Close()
var md struct {
AvailabilityZone string
}
if err := json.NewDecoder(r.Body).Decode(&md); err != nil {
return ""
}
// AWS MSK sets the `broker.rack` to AZ IDs rather than AZ names
return availabilityZoneMap[md.AvailabilityZone]
}