github.com/deanMdreon/kafka-go@v0.4.32/example_groupbalancer_test.go (about) 1 package kafka 2 3 import ( 4 "context" 5 "encoding/json" 6 "io/ioutil" 7 "net/http" 8 "os" 9 "strings" 10 "time" 11 ) 12 13 // ExampleNewReader_rackAffinity shows how the RackAffinityGroupBalancer can be 14 // used to pair up consumers with brokers in the same AWS availability zone. 15 // This code assumes that each brokers' rack is configured to be the name of the 16 // AZ in which it is running. 17 func ExampleNewReader_rackAffinity() { 18 r := NewReader(ReaderConfig{ 19 Brokers: []string{"kafka:9092"}, 20 GroupID: "my-group", 21 Topic: "my-topic", 22 GroupBalancers: []GroupBalancer{ 23 RackAffinityGroupBalancer{Rack: findRack()}, 24 RangeGroupBalancer{}, 25 }, 26 }) 27 28 r.ReadMessage(context.Background()) 29 30 r.Close() 31 } 32 33 // findRack is the basic rack resolver strategy for use in AWS. It supports 34 // * ECS with the task metadata endpoint enabled (returns the container 35 // instance's availability zone) 36 // * Linux EC2 (returns the instance's availability zone) 37 func findRack() string { 38 switch whereAmI() { 39 case "ecs": 40 return ecsAvailabilityZone() 41 case "ec2": 42 return ec2AvailabilityZone() 43 } 44 return "" 45 } 46 47 const ecsContainerMetadataURI = "ECS_CONTAINER_METADATA_URI" 48 49 // whereAmI determines which strategy the rack resolver should use. 50 func whereAmI() string { 51 // https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task-metadata-endpoint.html 52 if os.Getenv(ecsContainerMetadataURI) != "" { 53 return "ecs" 54 } 55 // https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/identify_ec2_instances.html 56 for _, path := range [...]string{ 57 "/sys/devices/virtual/dmi/id/product_uuid", 58 "/sys/hypervisor/uuid", 59 } { 60 b, err := ioutil.ReadFile(path) 61 if err != nil { 62 continue 63 } 64 s := string(b) 65 switch { 66 case strings.HasPrefix(s, "EC2"), strings.HasPrefix(s, "ec2"): 67 return "ec2" 68 } 69 } 70 return "somewhere" 71 } 72 73 // ecsAvailabilityZone queries the task endpoint for the metadata URI that ECS 74 // injects into the ECS_CONTAINER_METADATA_URI variable in order to retrieve 75 // the availability zone where the task is running. 76 func ecsAvailabilityZone() string { 77 client := http.Client{ 78 Timeout: time.Second, 79 Transport: &http.Transport{ 80 DisableCompression: true, 81 DisableKeepAlives: true, 82 }, 83 } 84 r, err := client.Get(os.Getenv(ecsContainerMetadataURI) + "/task") 85 if err != nil { 86 return "" 87 } 88 defer r.Body.Close() 89 90 var md struct { 91 AvailabilityZone string 92 } 93 if err := json.NewDecoder(r.Body).Decode(&md); err != nil { 94 return "" 95 } 96 return md.AvailabilityZone 97 } 98 99 // ec2AvailabilityZone queries the metadata endpoint to discover the 100 // availability zone where this code is running. we avoid calling this function 101 // unless we know we're in EC2. Otherwise, in other environments, we would need 102 // to wait for the request to 169.254.169.254 to timeout before proceeding. 103 func ec2AvailabilityZone() string { 104 client := http.Client{ 105 Timeout: time.Second, 106 Transport: &http.Transport{ 107 DisableCompression: true, 108 DisableKeepAlives: true, 109 }, 110 } 111 r, err := client.Get("http://169.254.169.254/latest/meta-data/placement/availability-zone") 112 if err != nil { 113 return "" 114 } 115 defer r.Body.Close() 116 b, err := ioutil.ReadAll(r.Body) 117 if err != nil { 118 return "" 119 } 120 return string(b) 121 }