github.com/QuangHoangHao/kafka-go@v0.4.36/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  }