add support for reading the first and last offsets of a topic/partition

ravi4422 · May 31, 2017 · cc7fc22 · cc7fc22
1 parent 9fd3208
commit cc7fc22
Show file tree

Hide file tree

Showing 6 changed files with 152 additions and 99 deletions.
diff --git a/conn.go b/conn.go
@@ -215,7 +215,6 @@ func (c *Conn) Seek(offset int64, whence int) (int64, error) {
 	}
 
 	first, last, err := c.ReadOffsets()
-	fmt.Println("first =", first, "last =", last)
 	if err != nil {
 		return 0, err
 	}
@@ -252,7 +251,19 @@ func (c *Conn) Read(b []byte) (int, error) {
 	var n int
 	err = c.readOperation(
 		func(id int32) error {
-			return nil
+			return c.writeRequest(fetchRequest, v1, id, fetchRequestV1{
+				ReplicaID:   -1,
+				MaxWaitTime: 100,
+				MinBytes:    1,
+				Topics: []fetchRequestTopicV1{{
+					TopicName: c.topic,
+					Partitions: []fetchRequestPartitionV1{{
+						Partition:   c.partition,
+						FetchOffset: offset,
+						MaxBytes:    int32(len(b) + 1),
+					}},
+				}},
+			})
 		},
 		func(size int) error {
 			return nil
@@ -302,7 +313,7 @@ func (c *Conn) ReadBatchAt(size int, offset int64) (*BatchReader, error) {
 
 	c.wlock.Lock()
 	c.conn.SetWriteDeadline(deadline)
-	err := c.writeRequest(fetchRequestKey, v1, id, nil) // TODO
+	err := c.writeRequest(fetchRequest, v1, id, nil) // TODO
 	c.wlock.Unlock()
 
 	if err != nil {
@@ -319,48 +330,57 @@ func (c *Conn) ReadBatchAt(size int, offset int64) (*BatchReader, error) {
 // ReadOffsets returns the absolute first and last offsets of the topic used by
 // the connection.
 func (c *Conn) ReadOffsets() (first int64, last int64, err error) {
-	first, last = -1, -1
+	type result struct {
+		off int64
+		err error
+	}
 
-	err = c.writeOperation(
-		func(id int32) error {
-			return c.writeRequest(offsetRequestKey, v1, id, listOffsetRequestV1{
-				ReplicaID: -1,
-				Topics: []listOffsetRequestTopicV1{{
-					TopicName: c.topic,
-					Partitions: []listOffsetRequestPartitionV1{
-						{Partition: c.partition, Time: -2},
-						{Partition: c.partition, Time: -1},
-					},
-				}},
-			})
-		},
-		func(size int) error {
-			var res []listOffsetResponseV1
-			if err := c.readResponse(size, &res); err != nil {
-				return err
-			}
-			fmt.Printf("%#v\n", res)
-			for _, r := range res {
-				for _, p := range r.PartitionOffsets {
-					if p.ErrorCode != 0 {
-						return Error(p.ErrorCode)
-					}
-					if first < 0 || p.Offset < first {
-						first = p.Offset
-					}
-					if last < 0 || p.Offset > last {
-						last = p.Offset
+	resch := make(chan result, 2)
+	fetch := func(time int64) {
+		off := int64(0)
+		err = c.writeOperation(
+			func(id int32) error {
+				return c.writeRequest(offsetRequest, v1, id, listOffsetRequestV1{
+					ReplicaID: -1,
+					Topics: []listOffsetRequestTopicV1{{
+						TopicName:  c.topic,
+						Partitions: []listOffsetRequestPartitionV1{{Partition: c.partition, Time: time}},
+					}},
+				})
+			},
+			func(size int) error {
+				var res []listOffsetResponseV1
+				if err := c.readResponse(size, &res); err != nil {
+					return err
+				}
+				for _, r := range res {
+					for _, p := range r.PartitionOffsets {
+						if p.ErrorCode != 0 {
+							return Error(p.ErrorCode)
+						} else {
+							off = p.Offset
+							return nil
+						}
 					}
 				}
-			}
-			return nil
-		},
-	)
-
-	if err == nil && (first < 0 || last < 0) {
-		err = UnknownTopicOrPartition
+				return UnknownTopicOrPartition
+			},
+		)
+		resch <- result{off, err}
 	}
 
+	// We have to submit two different requests to fetch the first and last
+	// offsets because kafka refuses requests that ask for multiple offsets
+	// on the same topic and partition.
+	go fetch(-1)
+	go fetch(-2)
+
+	res1 := <-resch
+	res2 := <-resch
+
+	first = minInt64(res1.off, res2.off)
+	last = maxInt64(res1.off, res2.off)
+	err = coalesceErrors(res1.err, res2.err)
 	return
 }
 
@@ -379,10 +399,10 @@ func (c *Conn) ReadPartitions(topics ...string) (partitions []Partition, err err
 
 	err = c.readOperation(
 		func(id int32) error {
-			return c.writeRequest(metadataRequestKey, v0, id, topicMetadataRequest(topics))
+			return c.writeRequest(metadataRequest, v0, id, topicMetadataRequestV0(topics))
 		},
 		func(size int) error {
-			var res metadataResponse
+			var res metadataResponseV0
 			if err := c.readResponse(size, &res); err != nil {
 				return err
 			}
@@ -448,12 +468,12 @@ func (c *Conn) Write(b []byte) (int, error) {
 
 	err := c.writeOperation(
 		func(id int32) error {
-			return c.writeRequest(produceRequestKey, v2, id, produceRequest{
+			return c.writeRequest(produceRequest, v2, id, produceRequestV2{
 				RequiredAcks: -1,
 				Timeout:      10000,
-				Topics: []produceRequestTopic{{
+				Topics: []produceRequestTopicV2{{
 					TopicName: c.topic,
-					Partitions: []produceRequestPartition{{
+					Partitions: []produceRequestPartitionV2{{
 						Partition:      c.partition,
 						MessageSetSize: sizeof(set),
 						MessageSet:     set,
@@ -462,7 +482,7 @@ func (c *Conn) Write(b []byte) (int, error) {
 			})
 		},
 		func(size int) error {
-			var res produceResponse
+			var res produceResponseV2
 			if err := c.readResponse(size, &res); err != nil {
 				return err
 			}

diff --git a/conn_test.go b/conn_test.go
@@ -10,6 +10,10 @@ import (
 	"time"
 )
 
+func init() {
+	rand.Seed(time.Now().UnixNano())
+}
+
 func TestConn(t *testing.T) {
 	tests := []struct {
 		scenario string
@@ -54,7 +58,7 @@ func TestConn(t *testing.T) {
 			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 			defer cancel()
 
-			topic := fmt.Sprintf("kafka-go-%d", rand.Int63())
+			topic := fmt.Sprintf("kafka-go-%016x", rand.Int63())
 
 			conn, err := (&Dialer{
 				Resolver: &net.Resolver{},

diff --git a/dialer.go b/dialer.go
@@ -135,7 +135,7 @@ func (d *Dialer) LookupLeader(ctx context.Context, network string, address strin
 	go func() {
 		for attempt := 0; true; attempt++ {
 			if attempt != 0 {
-				sleep(ctx, backoff(attempt, time.Second, time.Minute))
+				sleep(ctx, backoff(attempt, 100*time.Millisecond, 10*time.Second))
 			}
 
 			partitions, err := c.ReadPartitions(topic)

diff --git a/error.go b/error.go
@@ -337,3 +337,12 @@ func isTemporary(err error) bool {
 	})
 	return ok && e.Temporary()
 }
+
+func coalesceErrors(err ...error) error {
+	for _, e := range err {
+		if e != nil {
+			return e
+		}
+	}
+	return nil
+}