Merge pull request segmentio#212 from segmentio/brokers

Add Brokers function
huangfu · Feb 19, 2019 · fcb01ba · fcb01ba
2 parents 192634d + b5b8bef
commit fcb01ba
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/conn.go b/conn.go
@@ -149,7 +149,7 @@ func NewConnWith(conn net.Conn, config ConnConfig) *Conn {
 	return c
 }
 
-// GetController requests kafka for the current controller and returns its URL
+// Controller requests kafka for the current controller and returns its URL
 func (c *Conn) Controller() (broker Broker, err error) {
 	err = c.readOperation(
 		func(deadline time.Time, id int32) error {
@@ -176,6 +176,35 @@ func (c *Conn) Controller() (broker Broker, err error) {
 	return broker, err
 }
 
+// Brokers retrieve the broker list from the Kafka metadata
+func (c *Conn) Brokers() ([]Broker, error) {
+	var brokers []Broker
+	err := c.readOperation(
+		func(deadline time.Time, id int32) error {
+			return c.writeRequest(metadataRequest, v1, id, topicMetadataRequestV1([]string{}))
+		},
+		func(deadline time.Time, size int) error {
+			var res metadataResponseV1
+
+			if err := c.readResponse(size, &res); err != nil {
+				return err
+			}
+
+			brokers = make([]Broker, len(res.Brokers))
+			for i, brokerMeta := range res.Brokers {
+				brokers[i] = Broker{
+					ID:   int(brokerMeta.NodeID),
+					Port: int(brokerMeta.Port),
+					Host: brokerMeta.Host,
+					Rack: brokerMeta.Rack,
+				}
+			}
+			return nil
+		},
+	)
+	return brokers, err
+}
+
 // DeleteTopics deletes the specified topics.
 func (c *Conn) DeleteTopics(topics ...string) error {
 	_, err := c.deleteTopics(deleteTopicsRequestV0{

diff --git a/conn_test.go b/conn_test.go
@@ -245,6 +245,10 @@ func TestConn(t *testing.T) {
 			scenario: "test retrieve controller",
 			function: testController,
 		},
+		{
+			scenario: "test list brokers",
+			function: testBrokers,
+		},
 	}
 
 	const (
@@ -954,13 +958,28 @@ func testController(t *testing.T, conn *Conn) {
 		t.Errorf("expected 9092 received %d", b.Port)
 	}
 	if b.ID != 1 {
-		t.Errorf("expected 0 received %d", b.ID)
+		t.Errorf("expected 1 received %d", b.ID)
 	}
 	if b.Rack != "" {
 		t.Errorf("expected empty string for rack received %s", b.Rack)
 	}
 }
 
+func testBrokers(t *testing.T, conn *Conn) {
+	brokers, err := conn.Brokers()
+	if err != nil {
+		t.Error(err)
+	}
+
+	if len(brokers) != 1 {
+		t.Errorf("expected 1 broker in %+v", brokers)
+	}
+
+	if brokers[0].ID != 1 {
+		t.Errorf("expected ID 1 received %d", brokers[0].ID)
+	}
+}
+
 const benchmarkMessageCount = 100
 
 func BenchmarkConn(b *testing.B) {