No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.

1651 líneas
46 KiB

  1. package kafka
  2. import (
  3. "bufio"
  4. "errors"
  5. "fmt"
  6. "io"
  7. "math"
  8. "net"
  9. "os"
  10. "path/filepath"
  11. "sync"
  12. "sync/atomic"
  13. "time"
  14. )
  15. var (
  16. errInvalidWriteTopic = errors.New("writes must NOT set Topic on kafka.Message")
  17. errInvalidWritePartition = errors.New("writes must NOT set Partition on kafka.Message")
  18. )
  19. // Conn represents a connection to a kafka broker.
  20. //
  21. // Instances of Conn are safe to use concurrently from multiple goroutines.
  22. type Conn struct {
  23. // base network connection
  24. conn net.Conn
  25. // number of inflight requests on the connection.
  26. inflight int32
  27. // offset management (synchronized on the mutex field)
  28. mutex sync.Mutex
  29. offset int64
  30. // read buffer (synchronized on rlock)
  31. rlock sync.Mutex
  32. rbuf bufio.Reader
  33. // write buffer (synchronized on wlock)
  34. wlock sync.Mutex
  35. wbuf bufio.Writer
  36. wb writeBuffer
  37. // deadline management
  38. wdeadline connDeadline
  39. rdeadline connDeadline
  40. // immutable values of the connection object
  41. clientID string
  42. topic string
  43. partition int32
  44. fetchMaxBytes int32
  45. fetchMinSize int32
  46. broker int32
  47. rack string
  48. // correlation ID generator (synchronized on wlock)
  49. correlationID int32
  50. // number of replica acks required when publishing to a partition
  51. requiredAcks int32
  52. // lazily loaded API versions used by this connection
  53. apiVersions atomic.Value // apiVersionMap
  54. transactionalID *string
  55. }
  56. type apiVersionMap map[apiKey]ApiVersion
  57. func (v apiVersionMap) negotiate(key apiKey, sortedSupportedVersions ...apiVersion) apiVersion {
  58. x := v[key]
  59. for i := len(sortedSupportedVersions) - 1; i >= 0; i-- {
  60. s := sortedSupportedVersions[i]
  61. if apiVersion(x.MaxVersion) >= s {
  62. return s
  63. }
  64. }
  65. return -1
  66. }
  67. // ConnConfig is a configuration object used to create new instances of Conn.
  68. type ConnConfig struct {
  69. ClientID string
  70. Topic string
  71. Partition int
  72. Broker int
  73. Rack string
  74. // The transactional id to use for transactional delivery. Idempotent
  75. // deliver should be enabled if transactional id is configured.
  76. // For more details look at transactional.id description here: http://kafka.apache.org/documentation.html#producerconfigs
  77. // Empty string means that this connection can't be transactional.
  78. TransactionalID string
  79. }
  80. // ReadBatchConfig is a configuration object used for reading batches of messages.
  81. type ReadBatchConfig struct {
  82. // MinBytes indicates to the broker the minimum batch size that the consumer
  83. // will accept. Setting a high minimum when consuming from a low-volume topic
  84. // may result in delayed delivery when the broker does not have enough data to
  85. // satisfy the defined minimum.
  86. MinBytes int
  87. // MaxBytes indicates to the broker the maximum batch size that the consumer
  88. // will accept. The broker will truncate a message to satisfy this maximum, so
  89. // choose a value that is high enough for your largest message size.
  90. MaxBytes int
  91. // IsolationLevel controls the visibility of transactional records.
  92. // ReadUncommitted makes all records visible. With ReadCommitted only
  93. // non-transactional and committed records are visible.
  94. IsolationLevel IsolationLevel
  95. // MaxWait is the amount of time for the broker while waiting to hit the
  96. // min/max byte targets. This setting is independent of any network-level
  97. // timeouts or deadlines.
  98. //
  99. // For backward compatibility, when this field is left zero, kafka-go will
  100. // infer the max wait from the connection's read deadline.
  101. MaxWait time.Duration
  102. }
  103. type IsolationLevel int8
  104. const (
  105. ReadUncommitted IsolationLevel = 0
  106. ReadCommitted IsolationLevel = 1
  107. )
  108. var (
  109. // DefaultClientID is the default value used as ClientID of kafka
  110. // connections.
  111. DefaultClientID string
  112. )
  113. func init() {
  114. progname := filepath.Base(os.Args[0])
  115. hostname, _ := os.Hostname()
  116. DefaultClientID = fmt.Sprintf("%s@%s (github.com/segmentio/kafka-go)", progname, hostname)
  117. }
  118. // NewConn returns a new kafka connection for the given topic and partition.
  119. func NewConn(conn net.Conn, topic string, partition int) *Conn {
  120. return NewConnWith(conn, ConnConfig{
  121. Topic: topic,
  122. Partition: partition,
  123. })
  124. }
  125. func emptyToNullable(transactionalID string) (result *string) {
  126. if transactionalID != "" {
  127. result = &transactionalID
  128. }
  129. return result
  130. }
  131. // NewConnWith returns a new kafka connection configured with config.
  132. // The offset is initialized to FirstOffset.
  133. func NewConnWith(conn net.Conn, config ConnConfig) *Conn {
  134. if len(config.ClientID) == 0 {
  135. config.ClientID = DefaultClientID
  136. }
  137. if config.Partition < 0 || config.Partition > math.MaxInt32 {
  138. panic(fmt.Sprintf("invalid partition number: %d", config.Partition))
  139. }
  140. c := &Conn{
  141. conn: conn,
  142. rbuf: *bufio.NewReader(conn),
  143. wbuf: *bufio.NewWriter(conn),
  144. clientID: config.ClientID,
  145. topic: config.Topic,
  146. partition: int32(config.Partition),
  147. broker: int32(config.Broker),
  148. rack: config.Rack,
  149. offset: FirstOffset,
  150. requiredAcks: -1,
  151. transactionalID: emptyToNullable(config.TransactionalID),
  152. }
  153. c.wb.w = &c.wbuf
  154. // The fetch request needs to ask for a MaxBytes value that is at least
  155. // enough to load the control data of the response. To avoid having to
  156. // recompute it on every read, it is cached here in the Conn value.
  157. c.fetchMinSize = (fetchResponseV2{
  158. Topics: []fetchResponseTopicV2{{
  159. TopicName: config.Topic,
  160. Partitions: []fetchResponsePartitionV2{{
  161. Partition: int32(config.Partition),
  162. MessageSet: messageSet{{}},
  163. }},
  164. }},
  165. }).size()
  166. c.fetchMaxBytes = math.MaxInt32 - c.fetchMinSize
  167. return c
  168. }
  169. func (c *Conn) negotiateVersion(key apiKey, sortedSupportedVersions ...apiVersion) (apiVersion, error) {
  170. v, err := c.loadVersions()
  171. if err != nil {
  172. return -1, err
  173. }
  174. a := v.negotiate(key, sortedSupportedVersions...)
  175. if a < 0 {
  176. return -1, fmt.Errorf("no matching versions were found between the client and the broker for API key %d", key)
  177. }
  178. return a, nil
  179. }
  180. func (c *Conn) loadVersions() (apiVersionMap, error) {
  181. v, _ := c.apiVersions.Load().(apiVersionMap)
  182. if v != nil {
  183. return v, nil
  184. }
  185. brokerVersions, err := c.ApiVersions()
  186. if err != nil {
  187. return nil, err
  188. }
  189. v = make(apiVersionMap, len(brokerVersions))
  190. for _, a := range brokerVersions {
  191. v[apiKey(a.ApiKey)] = a
  192. }
  193. c.apiVersions.Store(v)
  194. return v, nil
  195. }
  196. // Broker returns a Broker value representing the kafka broker that this
  197. // connection was established to.
  198. func (c *Conn) Broker() Broker {
  199. addr := c.conn.RemoteAddr()
  200. host, port, _ := splitHostPortNumber(addr.String())
  201. return Broker{
  202. Host: host,
  203. Port: port,
  204. ID: int(c.broker),
  205. Rack: c.rack,
  206. }
  207. }
  208. // Controller requests kafka for the current controller and returns its URL.
  209. func (c *Conn) Controller() (broker Broker, err error) {
  210. err = c.readOperation(
  211. func(deadline time.Time, id int32) error {
  212. return c.writeRequest(metadata, v1, id, topicMetadataRequestV1([]string{}))
  213. },
  214. func(deadline time.Time, size int) error {
  215. var res metadataResponseV1
  216. if err := c.readResponse(size, &res); err != nil {
  217. return err
  218. }
  219. for _, brokerMeta := range res.Brokers {
  220. if brokerMeta.NodeID == res.ControllerID {
  221. broker = Broker{ID: int(brokerMeta.NodeID),
  222. Port: int(brokerMeta.Port),
  223. Host: brokerMeta.Host,
  224. Rack: brokerMeta.Rack}
  225. break
  226. }
  227. }
  228. return nil
  229. },
  230. )
  231. return broker, err
  232. }
  233. // Brokers retrieve the broker list from the Kafka metadata.
  234. func (c *Conn) Brokers() ([]Broker, error) {
  235. var brokers []Broker
  236. err := c.readOperation(
  237. func(deadline time.Time, id int32) error {
  238. return c.writeRequest(metadata, v1, id, topicMetadataRequestV1([]string{}))
  239. },
  240. func(deadline time.Time, size int) error {
  241. var res metadataResponseV1
  242. if err := c.readResponse(size, &res); err != nil {
  243. return err
  244. }
  245. brokers = make([]Broker, len(res.Brokers))
  246. for i, brokerMeta := range res.Brokers {
  247. brokers[i] = Broker{
  248. ID: int(brokerMeta.NodeID),
  249. Port: int(brokerMeta.Port),
  250. Host: brokerMeta.Host,
  251. Rack: brokerMeta.Rack,
  252. }
  253. }
  254. return nil
  255. },
  256. )
  257. return brokers, err
  258. }
  259. // DeleteTopics deletes the specified topics.
  260. func (c *Conn) DeleteTopics(topics ...string) error {
  261. _, err := c.deleteTopics(deleteTopicsRequest{
  262. Topics: topics,
  263. })
  264. return err
  265. }
  266. // findCoordinator finds the coordinator for the specified group or transaction
  267. //
  268. // See http://kafka.apache.org/protocol.html#The_Messages_FindCoordinator
  269. func (c *Conn) findCoordinator(request findCoordinatorRequestV0) (findCoordinatorResponseV0, error) {
  270. var response findCoordinatorResponseV0
  271. err := c.readOperation(
  272. func(deadline time.Time, id int32) error {
  273. return c.writeRequest(findCoordinator, v0, id, request)
  274. },
  275. func(deadline time.Time, size int) error {
  276. return expectZeroSize(func() (remain int, err error) {
  277. return (&response).readFrom(&c.rbuf, size)
  278. }())
  279. },
  280. )
  281. if err != nil {
  282. return findCoordinatorResponseV0{}, err
  283. }
  284. if response.ErrorCode != 0 {
  285. return findCoordinatorResponseV0{}, Error(response.ErrorCode)
  286. }
  287. return response, nil
  288. }
  289. // heartbeat sends a heartbeat message required by consumer groups
  290. //
  291. // See http://kafka.apache.org/protocol.html#The_Messages_Heartbeat
  292. func (c *Conn) heartbeat(request heartbeatRequestV0) (heartbeatResponseV0, error) {
  293. var response heartbeatResponseV0
  294. err := c.writeOperation(
  295. func(deadline time.Time, id int32) error {
  296. return c.writeRequest(heartbeat, v0, id, request)
  297. },
  298. func(deadline time.Time, size int) error {
  299. return expectZeroSize(func() (remain int, err error) {
  300. return (&response).readFrom(&c.rbuf, size)
  301. }())
  302. },
  303. )
  304. if err != nil {
  305. return heartbeatResponseV0{}, err
  306. }
  307. if response.ErrorCode != 0 {
  308. return heartbeatResponseV0{}, Error(response.ErrorCode)
  309. }
  310. return response, nil
  311. }
  312. // joinGroup attempts to join a consumer group
  313. //
  314. // See http://kafka.apache.org/protocol.html#The_Messages_JoinGroup
  315. func (c *Conn) joinGroup(request joinGroupRequest) (joinGroupResponse, error) {
  316. version, err := c.negotiateVersion(joinGroup, v1, v2)
  317. if err != nil {
  318. return joinGroupResponse{}, err
  319. }
  320. response := joinGroupResponse{v: version}
  321. err = c.writeOperation(
  322. func(deadline time.Time, id int32) error {
  323. return c.writeRequest(joinGroup, version, id, request)
  324. },
  325. func(deadline time.Time, size int) error {
  326. return expectZeroSize(func() (remain int, err error) {
  327. return (&response).readFrom(&c.rbuf, size)
  328. }())
  329. },
  330. )
  331. if err != nil {
  332. return joinGroupResponse{}, err
  333. }
  334. if response.ErrorCode != 0 {
  335. return joinGroupResponse{}, Error(response.ErrorCode)
  336. }
  337. return response, nil
  338. }
  339. // leaveGroup leaves the consumer from the consumer group
  340. //
  341. // See http://kafka.apache.org/protocol.html#The_Messages_LeaveGroup
  342. func (c *Conn) leaveGroup(request leaveGroupRequestV0) (leaveGroupResponseV0, error) {
  343. var response leaveGroupResponseV0
  344. err := c.writeOperation(
  345. func(deadline time.Time, id int32) error {
  346. return c.writeRequest(leaveGroup, v0, id, request)
  347. },
  348. func(deadline time.Time, size int) error {
  349. return expectZeroSize(func() (remain int, err error) {
  350. return (&response).readFrom(&c.rbuf, size)
  351. }())
  352. },
  353. )
  354. if err != nil {
  355. return leaveGroupResponseV0{}, err
  356. }
  357. if response.ErrorCode != 0 {
  358. return leaveGroupResponseV0{}, Error(response.ErrorCode)
  359. }
  360. return response, nil
  361. }
  362. // listGroups lists all the consumer groups
  363. //
  364. // See http://kafka.apache.org/protocol.html#The_Messages_ListGroups
  365. func (c *Conn) listGroups(request listGroupsRequestV1) (listGroupsResponseV1, error) {
  366. var response listGroupsResponseV1
  367. err := c.readOperation(
  368. func(deadline time.Time, id int32) error {
  369. return c.writeRequest(listGroups, v1, id, request)
  370. },
  371. func(deadline time.Time, size int) error {
  372. return expectZeroSize(func() (remain int, err error) {
  373. return (&response).readFrom(&c.rbuf, size)
  374. }())
  375. },
  376. )
  377. if err != nil {
  378. return listGroupsResponseV1{}, err
  379. }
  380. if response.ErrorCode != 0 {
  381. return listGroupsResponseV1{}, Error(response.ErrorCode)
  382. }
  383. return response, nil
  384. }
  385. // offsetCommit commits the specified topic partition offsets
  386. //
  387. // See http://kafka.apache.org/protocol.html#The_Messages_OffsetCommit
  388. func (c *Conn) offsetCommit(request offsetCommitRequestV2) (offsetCommitResponseV2, error) {
  389. var response offsetCommitResponseV2
  390. err := c.writeOperation(
  391. func(deadline time.Time, id int32) error {
  392. return c.writeRequest(offsetCommit, v2, id, request)
  393. },
  394. func(deadline time.Time, size int) error {
  395. return expectZeroSize(func() (remain int, err error) {
  396. return (&response).readFrom(&c.rbuf, size)
  397. }())
  398. },
  399. )
  400. if err != nil {
  401. return offsetCommitResponseV2{}, err
  402. }
  403. for _, r := range response.Responses {
  404. for _, pr := range r.PartitionResponses {
  405. if pr.ErrorCode != 0 {
  406. return offsetCommitResponseV2{}, Error(pr.ErrorCode)
  407. }
  408. }
  409. }
  410. return response, nil
  411. }
  412. // offsetFetch fetches the offsets for the specified topic partitions.
  413. // -1 indicates that there is no offset saved for the partition.
  414. //
  415. // See http://kafka.apache.org/protocol.html#The_Messages_OffsetFetch
  416. func (c *Conn) offsetFetch(request offsetFetchRequestV1) (offsetFetchResponseV1, error) {
  417. var response offsetFetchResponseV1
  418. err := c.readOperation(
  419. func(deadline time.Time, id int32) error {
  420. return c.writeRequest(offsetFetch, v1, id, request)
  421. },
  422. func(deadline time.Time, size int) error {
  423. return expectZeroSize(func() (remain int, err error) {
  424. return (&response).readFrom(&c.rbuf, size)
  425. }())
  426. },
  427. )
  428. if err != nil {
  429. return offsetFetchResponseV1{}, err
  430. }
  431. for _, r := range response.Responses {
  432. for _, pr := range r.PartitionResponses {
  433. if pr.ErrorCode != 0 {
  434. return offsetFetchResponseV1{}, Error(pr.ErrorCode)
  435. }
  436. }
  437. }
  438. return response, nil
  439. }
  440. // syncGroup completes the handshake to join a consumer group
  441. //
  442. // See http://kafka.apache.org/protocol.html#The_Messages_SyncGroup
  443. func (c *Conn) syncGroup(request syncGroupRequestV0) (syncGroupResponseV0, error) {
  444. var response syncGroupResponseV0
  445. err := c.readOperation(
  446. func(deadline time.Time, id int32) error {
  447. return c.writeRequest(syncGroup, v0, id, request)
  448. },
  449. func(deadline time.Time, size int) error {
  450. return expectZeroSize(func() (remain int, err error) {
  451. return (&response).readFrom(&c.rbuf, size)
  452. }())
  453. },
  454. )
  455. if err != nil {
  456. return syncGroupResponseV0{}, err
  457. }
  458. if response.ErrorCode != 0 {
  459. return syncGroupResponseV0{}, Error(response.ErrorCode)
  460. }
  461. return response, nil
  462. }
  463. // Close closes the kafka connection.
  464. func (c *Conn) Close() error {
  465. return c.conn.Close()
  466. }
  467. // LocalAddr returns the local network address.
  468. func (c *Conn) LocalAddr() net.Addr {
  469. return c.conn.LocalAddr()
  470. }
  471. // RemoteAddr returns the remote network address.
  472. func (c *Conn) RemoteAddr() net.Addr {
  473. return c.conn.RemoteAddr()
  474. }
  475. // SetDeadline sets the read and write deadlines associated with the connection.
  476. // It is equivalent to calling both SetReadDeadline and SetWriteDeadline.
  477. //
  478. // A deadline is an absolute time after which I/O operations fail with a timeout
  479. // (see type Error) instead of blocking. The deadline applies to all future and
  480. // pending I/O, not just the immediately following call to Read or Write. After
  481. // a deadline has been exceeded, the connection may be closed if it was found to
  482. // be in an unrecoverable state.
  483. //
  484. // A zero value for t means I/O operations will not time out.
  485. func (c *Conn) SetDeadline(t time.Time) error {
  486. c.rdeadline.setDeadline(t)
  487. c.wdeadline.setDeadline(t)
  488. return nil
  489. }
  490. // SetReadDeadline sets the deadline for future Read calls and any
  491. // currently-blocked Read call.
  492. // A zero value for t means Read will not time out.
  493. func (c *Conn) SetReadDeadline(t time.Time) error {
  494. c.rdeadline.setDeadline(t)
  495. return nil
  496. }
  497. // SetWriteDeadline sets the deadline for future Write calls and any
  498. // currently-blocked Write call.
  499. // Even if write times out, it may return n > 0, indicating that some of the
  500. // data was successfully written.
  501. // A zero value for t means Write will not time out.
  502. func (c *Conn) SetWriteDeadline(t time.Time) error {
  503. c.wdeadline.setDeadline(t)
  504. return nil
  505. }
  506. // Offset returns the current offset of the connection as pair of integers,
  507. // where the first one is an offset value and the second one indicates how
  508. // to interpret it.
  509. //
  510. // See Seek for more details about the offset and whence values.
  511. func (c *Conn) Offset() (offset int64, whence int) {
  512. c.mutex.Lock()
  513. offset = c.offset
  514. c.mutex.Unlock()
  515. switch offset {
  516. case FirstOffset:
  517. offset = 0
  518. whence = SeekStart
  519. case LastOffset:
  520. offset = 0
  521. whence = SeekEnd
  522. default:
  523. whence = SeekAbsolute
  524. }
  525. return
  526. }
  527. const (
  528. SeekStart = 0 // Seek relative to the first offset available in the partition.
  529. SeekAbsolute = 1 // Seek to an absolute offset.
  530. SeekEnd = 2 // Seek relative to the last offset available in the partition.
  531. SeekCurrent = 3 // Seek relative to the current offset.
  532. // This flag may be combined to any of the SeekAbsolute and SeekCurrent
  533. // constants to skip the bound check that the connection would do otherwise.
  534. // Programs can use this flag to avoid making a metadata request to the kafka
  535. // broker to read the current first and last offsets of the partition.
  536. SeekDontCheck = 1 << 30
  537. )
  538. // Seek sets the offset for the next read or write operation according to whence, which
  539. // should be one of SeekStart, SeekAbsolute, SeekEnd, or SeekCurrent.
  540. // When seeking relative to the end, the offset is subtracted from the current offset.
  541. // Note that for historical reasons, these do not align with the usual whence constants
  542. // as in lseek(2) or os.Seek.
  543. // The method returns the new absolute offset of the connection.
  544. func (c *Conn) Seek(offset int64, whence int) (int64, error) {
  545. seekDontCheck := (whence & SeekDontCheck) != 0
  546. whence &= ^SeekDontCheck
  547. switch whence {
  548. case SeekStart, SeekAbsolute, SeekEnd, SeekCurrent:
  549. default:
  550. return 0, fmt.Errorf("whence must be one of 0, 1, 2, or 3. (whence = %d)", whence)
  551. }
  552. if seekDontCheck {
  553. if whence == SeekAbsolute {
  554. c.mutex.Lock()
  555. c.offset = offset
  556. c.mutex.Unlock()
  557. return offset, nil
  558. }
  559. if whence == SeekCurrent {
  560. c.mutex.Lock()
  561. c.offset += offset
  562. offset = c.offset
  563. c.mutex.Unlock()
  564. return offset, nil
  565. }
  566. }
  567. if whence == SeekAbsolute {
  568. c.mutex.Lock()
  569. unchanged := offset == c.offset
  570. c.mutex.Unlock()
  571. if unchanged {
  572. return offset, nil
  573. }
  574. }
  575. if whence == SeekCurrent {
  576. c.mutex.Lock()
  577. offset = c.offset + offset
  578. c.mutex.Unlock()
  579. }
  580. first, last, err := c.ReadOffsets()
  581. if err != nil {
  582. return 0, err
  583. }
  584. switch whence {
  585. case SeekStart:
  586. offset = first + offset
  587. case SeekEnd:
  588. offset = last - offset
  589. }
  590. if offset < first || offset > last {
  591. return 0, OffsetOutOfRange
  592. }
  593. c.mutex.Lock()
  594. c.offset = offset
  595. c.mutex.Unlock()
  596. return offset, nil
  597. }
  598. // Read reads the message at the current offset from the connection, advancing
  599. // the offset on success so the next call to a read method will produce the next
  600. // message.
  601. // The method returns the number of bytes read, or an error if something went
  602. // wrong.
  603. //
  604. // While it is safe to call Read concurrently from multiple goroutines it may
  605. // be hard for the program to predict the results as the connection offset will
  606. // be read and written by multiple goroutines, they could read duplicates, or
  607. // messages may be seen by only some of the goroutines.
  608. //
  609. // The method fails with io.ErrShortBuffer if the buffer passed as argument is
  610. // too small to hold the message value.
  611. //
  612. // This method is provided to satisfy the net.Conn interface but is much less
  613. // efficient than using the more general purpose ReadBatch method.
  614. func (c *Conn) Read(b []byte) (int, error) {
  615. batch := c.ReadBatch(1, len(b))
  616. n, err := batch.Read(b)
  617. return n, coalesceErrors(silentEOF(err), batch.Close())
  618. }
  619. // ReadMessage reads the message at the current offset from the connection,
  620. // advancing the offset on success so the next call to a read method will
  621. // produce the next message.
  622. //
  623. // Because this method allocate memory buffers for the message key and value
  624. // it is less memory-efficient than Read, but has the advantage of never
  625. // failing with io.ErrShortBuffer.
  626. //
  627. // While it is safe to call Read concurrently from multiple goroutines it may
  628. // be hard for the program to predict the results as the connection offset will
  629. // be read and written by multiple goroutines, they could read duplicates, or
  630. // messages may be seen by only some of the goroutines.
  631. //
  632. // This method is provided for convenience purposes but is much less efficient
  633. // than using the more general purpose ReadBatch method.
  634. func (c *Conn) ReadMessage(maxBytes int) (Message, error) {
  635. batch := c.ReadBatch(1, maxBytes)
  636. msg, err := batch.ReadMessage()
  637. return msg, coalesceErrors(silentEOF(err), batch.Close())
  638. }
  639. // ReadBatch reads a batch of messages from the kafka server. The method always
  640. // returns a non-nil Batch value. If an error occurred, either sending the fetch
  641. // request or reading the response, the error will be made available by the
  642. // returned value of the batch's Close method.
  643. //
  644. // While it is safe to call ReadBatch concurrently from multiple goroutines it
  645. // may be hard for the program to predict the results as the connection offset
  646. // will be read and written by multiple goroutines, they could read duplicates,
  647. // or messages may be seen by only some of the goroutines.
  648. //
  649. // A program doesn't specify the number of messages in wants from a batch, but
  650. // gives the minimum and maximum number of bytes that it wants to receive from
  651. // the kafka server.
  652. func (c *Conn) ReadBatch(minBytes, maxBytes int) *Batch {
  653. return c.ReadBatchWith(ReadBatchConfig{
  654. MinBytes: minBytes,
  655. MaxBytes: maxBytes,
  656. })
  657. }
  658. // ReadBatchWith in every way is similar to ReadBatch. ReadBatch is configured
  659. // with the default values in ReadBatchConfig except for minBytes and maxBytes.
  660. func (c *Conn) ReadBatchWith(cfg ReadBatchConfig) *Batch {
  661. var adjustedDeadline time.Time
  662. var maxFetch = int(c.fetchMaxBytes)
  663. if cfg.MinBytes < 0 || cfg.MinBytes > maxFetch {
  664. return &Batch{err: fmt.Errorf("kafka.(*Conn).ReadBatch: minBytes of %d out of [1,%d] bounds", cfg.MinBytes, maxFetch)}
  665. }
  666. if cfg.MaxBytes < 0 || cfg.MaxBytes > maxFetch {
  667. return &Batch{err: fmt.Errorf("kafka.(*Conn).ReadBatch: maxBytes of %d out of [1,%d] bounds", cfg.MaxBytes, maxFetch)}
  668. }
  669. if cfg.MinBytes > cfg.MaxBytes {
  670. return &Batch{err: fmt.Errorf("kafka.(*Conn).ReadBatch: minBytes (%d) > maxBytes (%d)", cfg.MinBytes, cfg.MaxBytes)}
  671. }
  672. offset, whence := c.Offset()
  673. offset, err := c.Seek(offset, whence|SeekDontCheck)
  674. if err != nil {
  675. return &Batch{err: dontExpectEOF(err)}
  676. }
  677. fetchVersion, err := c.negotiateVersion(fetch, v2, v5, v10)
  678. if err != nil {
  679. return &Batch{err: dontExpectEOF(err)}
  680. }
  681. id, err := c.doRequest(&c.rdeadline, func(deadline time.Time, id int32) error {
  682. now := time.Now()
  683. var timeout time.Duration
  684. if cfg.MaxWait > 0 {
  685. // explicitly-configured case: no changes are made to the deadline,
  686. // and the timeout is sent exactly as specified.
  687. timeout = cfg.MaxWait
  688. } else {
  689. // default case: use the original logic to adjust the conn's
  690. // deadline.T
  691. deadline = adjustDeadlineForRTT(deadline, now, defaultRTT)
  692. timeout = deadlineToTimeout(deadline, now)
  693. }
  694. // save this variable outside of the closure for later use in detecting
  695. // truncated messages.
  696. adjustedDeadline = deadline
  697. switch fetchVersion {
  698. case v10:
  699. return c.wb.writeFetchRequestV10(
  700. id,
  701. c.clientID,
  702. c.topic,
  703. c.partition,
  704. offset,
  705. cfg.MinBytes,
  706. cfg.MaxBytes+int(c.fetchMinSize),
  707. timeout,
  708. int8(cfg.IsolationLevel),
  709. )
  710. case v5:
  711. return c.wb.writeFetchRequestV5(
  712. id,
  713. c.clientID,
  714. c.topic,
  715. c.partition,
  716. offset,
  717. cfg.MinBytes,
  718. cfg.MaxBytes+int(c.fetchMinSize),
  719. timeout,
  720. int8(cfg.IsolationLevel),
  721. )
  722. default:
  723. return c.wb.writeFetchRequestV2(
  724. id,
  725. c.clientID,
  726. c.topic,
  727. c.partition,
  728. offset,
  729. cfg.MinBytes,
  730. cfg.MaxBytes+int(c.fetchMinSize),
  731. timeout,
  732. )
  733. }
  734. })
  735. if err != nil {
  736. return &Batch{err: dontExpectEOF(err)}
  737. }
  738. _, size, lock, err := c.waitResponse(&c.rdeadline, id)
  739. if err != nil {
  740. return &Batch{err: dontExpectEOF(err)}
  741. }
  742. var throttle int32
  743. var highWaterMark int64
  744. var remain int
  745. switch fetchVersion {
  746. case v10:
  747. throttle, highWaterMark, remain, err = readFetchResponseHeaderV10(&c.rbuf, size)
  748. case v5:
  749. throttle, highWaterMark, remain, err = readFetchResponseHeaderV5(&c.rbuf, size)
  750. default:
  751. throttle, highWaterMark, remain, err = readFetchResponseHeaderV2(&c.rbuf, size)
  752. }
  753. if errors.Is(err, errShortRead) {
  754. err = checkTimeoutErr(adjustedDeadline)
  755. }
  756. var msgs *messageSetReader
  757. if err == nil {
  758. if highWaterMark == offset {
  759. msgs = &messageSetReader{empty: true}
  760. } else {
  761. msgs, err = newMessageSetReader(&c.rbuf, remain)
  762. }
  763. }
  764. if errors.Is(err, errShortRead) {
  765. err = checkTimeoutErr(adjustedDeadline)
  766. }
  767. return &Batch{
  768. conn: c,
  769. msgs: msgs,
  770. deadline: adjustedDeadline,
  771. throttle: makeDuration(throttle),
  772. lock: lock,
  773. topic: c.topic, // topic is copied to Batch to prevent race with Batch.close
  774. partition: int(c.partition), // partition is copied to Batch to prevent race with Batch.close
  775. offset: offset,
  776. highWaterMark: highWaterMark,
  777. // there shouldn't be a short read on initially setting up the batch.
  778. // as such, any io.EOF is re-mapped to an io.ErrUnexpectedEOF so that we
  779. // don't accidentally signal that we successfully reached the end of the
  780. // batch.
  781. err: dontExpectEOF(err),
  782. }
  783. }
  784. // ReadOffset returns the offset of the first message with a timestamp equal or
  785. // greater to t.
  786. func (c *Conn) ReadOffset(t time.Time) (int64, error) {
  787. return c.readOffset(timestamp(t))
  788. }
  789. // ReadFirstOffset returns the first offset available on the connection.
  790. func (c *Conn) ReadFirstOffset() (int64, error) {
  791. return c.readOffset(FirstOffset)
  792. }
  793. // ReadLastOffset returns the last offset available on the connection.
  794. func (c *Conn) ReadLastOffset() (int64, error) {
  795. return c.readOffset(LastOffset)
  796. }
  797. // ReadOffsets returns the absolute first and last offsets of the topic used by
  798. // the connection.
  799. func (c *Conn) ReadOffsets() (first, last int64, err error) {
  800. // We have to submit two different requests to fetch the first and last
  801. // offsets because kafka refuses requests that ask for multiple offsets
  802. // on the same topic and partition.
  803. if first, err = c.ReadFirstOffset(); err != nil {
  804. return
  805. }
  806. if last, err = c.ReadLastOffset(); err != nil {
  807. first = 0 // don't leak the value on error
  808. return
  809. }
  810. return
  811. }
  812. func (c *Conn) readOffset(t int64) (offset int64, err error) {
  813. err = c.readOperation(
  814. func(deadline time.Time, id int32) error {
  815. return c.wb.writeListOffsetRequestV1(id, c.clientID, c.topic, c.partition, t)
  816. },
  817. func(deadline time.Time, size int) error {
  818. return expectZeroSize(readArrayWith(&c.rbuf, size, func(r *bufio.Reader, size int) (int, error) {
  819. // We skip the topic name because we've made a request for
  820. // a single topic.
  821. size, err := discardString(r, size)
  822. if err != nil {
  823. return size, err
  824. }
  825. // Reading the array of partitions, there will be only one
  826. // partition which gives the offset we're looking for.
  827. return readArrayWith(r, size, func(r *bufio.Reader, size int) (int, error) {
  828. var p partitionOffsetV1
  829. size, err := p.readFrom(r, size)
  830. if err != nil {
  831. return size, err
  832. }
  833. if p.ErrorCode != 0 {
  834. return size, Error(p.ErrorCode)
  835. }
  836. offset = p.Offset
  837. return size, nil
  838. })
  839. }))
  840. },
  841. )
  842. return
  843. }
  844. // ReadPartitions returns the list of available partitions for the given list of
  845. // topics.
  846. //
  847. // If the method is called with no topic, it uses the topic configured on the
  848. // connection. If there are none, the method fetches all partitions of the kafka
  849. // cluster.
  850. func (c *Conn) ReadPartitions(topics ...string) (partitions []Partition, err error) {
  851. if len(topics) == 0 {
  852. if len(c.topic) != 0 {
  853. defaultTopics := [...]string{c.topic}
  854. topics = defaultTopics[:]
  855. } else {
  856. // topics needs to be explicitly nil-ed out or the broker will
  857. // interpret it as a request for 0 partitions instead of all.
  858. topics = nil
  859. }
  860. }
  861. metadataVersion, err := c.negotiateVersion(metadata, v1, v6)
  862. if err != nil {
  863. return nil, err
  864. }
  865. err = c.readOperation(
  866. func(deadline time.Time, id int32) error {
  867. switch metadataVersion {
  868. case v6:
  869. return c.writeRequest(metadata, v6, id, topicMetadataRequestV6{Topics: topics, AllowAutoTopicCreation: true})
  870. default:
  871. return c.writeRequest(metadata, v1, id, topicMetadataRequestV1(topics))
  872. }
  873. },
  874. func(deadline time.Time, size int) error {
  875. partitions, err = c.readPartitionsResponse(metadataVersion, size)
  876. return err
  877. },
  878. )
  879. return
  880. }
  881. func (c *Conn) readPartitionsResponse(metadataVersion apiVersion, size int) ([]Partition, error) {
  882. switch metadataVersion {
  883. case v6:
  884. var res metadataResponseV6
  885. if err := c.readResponse(size, &res); err != nil {
  886. return nil, err
  887. }
  888. brokers := readBrokerMetadata(res.Brokers)
  889. return c.readTopicMetadatav6(brokers, res.Topics)
  890. default:
  891. var res metadataResponseV1
  892. if err := c.readResponse(size, &res); err != nil {
  893. return nil, err
  894. }
  895. brokers := readBrokerMetadata(res.Brokers)
  896. return c.readTopicMetadatav1(brokers, res.Topics)
  897. }
  898. }
  899. func readBrokerMetadata(brokerMetadata []brokerMetadataV1) map[int32]Broker {
  900. brokers := make(map[int32]Broker, len(brokerMetadata))
  901. for _, b := range brokerMetadata {
  902. brokers[b.NodeID] = Broker{
  903. Host: b.Host,
  904. Port: int(b.Port),
  905. ID: int(b.NodeID),
  906. Rack: b.Rack,
  907. }
  908. }
  909. return brokers
  910. }
  911. func (c *Conn) readTopicMetadatav1(brokers map[int32]Broker, topicMetadata []topicMetadataV1) (partitions []Partition, err error) {
  912. for _, t := range topicMetadata {
  913. if t.TopicErrorCode != 0 && (c.topic == "" || t.TopicName == c.topic) {
  914. // We only report errors if they happened for the topic of
  915. // the connection, otherwise the topic will simply have no
  916. // partitions in the result set.
  917. return nil, Error(t.TopicErrorCode)
  918. }
  919. for _, p := range t.Partitions {
  920. partitions = append(partitions, Partition{
  921. Topic: t.TopicName,
  922. Leader: brokers[p.Leader],
  923. Replicas: makeBrokers(brokers, p.Replicas...),
  924. Isr: makeBrokers(brokers, p.Isr...),
  925. ID: int(p.PartitionID),
  926. OfflineReplicas: []Broker{},
  927. })
  928. }
  929. }
  930. return
  931. }
  932. func (c *Conn) readTopicMetadatav6(brokers map[int32]Broker, topicMetadata []topicMetadataV6) (partitions []Partition, err error) {
  933. for _, t := range topicMetadata {
  934. if t.TopicErrorCode != 0 && (c.topic == "" || t.TopicName == c.topic) {
  935. // We only report errors if they happened for the topic of
  936. // the connection, otherwise the topic will simply have no
  937. // partitions in the result set.
  938. return nil, Error(t.TopicErrorCode)
  939. }
  940. for _, p := range t.Partitions {
  941. partitions = append(partitions, Partition{
  942. Topic: t.TopicName,
  943. Leader: brokers[p.Leader],
  944. Replicas: makeBrokers(brokers, p.Replicas...),
  945. Isr: makeBrokers(brokers, p.Isr...),
  946. ID: int(p.PartitionID),
  947. OfflineReplicas: makeBrokers(brokers, p.OfflineReplicas...),
  948. })
  949. }
  950. }
  951. return
  952. }
  953. func makeBrokers(brokers map[int32]Broker, ids ...int32) []Broker {
  954. b := make([]Broker, len(ids))
  955. for i, id := range ids {
  956. br, ok := brokers[id]
  957. if !ok {
  958. // When the broker id isn't found in the current list of known
  959. // brokers, use a placeholder to report that the cluster has
  960. // logical knowledge of the broker but no information about the
  961. // physical host where it is running.
  962. br.ID = int(id)
  963. }
  964. b[i] = br
  965. }
  966. return b
  967. }
  968. // Write writes a message to the kafka broker that this connection was
  969. // established to. The method returns the number of bytes written, or an error
  970. // if something went wrong.
  971. //
  972. // The operation either succeeds or fail, it never partially writes the message.
  973. //
  974. // This method is exposed to satisfy the net.Conn interface but is less efficient
  975. // than the more general purpose WriteMessages method.
  976. func (c *Conn) Write(b []byte) (int, error) {
  977. return c.WriteCompressedMessages(nil, Message{Value: b})
  978. }
  979. // WriteMessages writes a batch of messages to the connection's topic and
  980. // partition, returning the number of bytes written. The write is an atomic
  981. // operation, it either fully succeeds or fails.
  982. func (c *Conn) WriteMessages(msgs ...Message) (int, error) {
  983. return c.WriteCompressedMessages(nil, msgs...)
  984. }
  985. // WriteCompressedMessages writes a batch of messages to the connection's topic
  986. // and partition, returning the number of bytes written. The write is an atomic
  987. // operation, it either fully succeeds or fails.
  988. //
  989. // If the compression codec is not nil, the messages will be compressed.
  990. func (c *Conn) WriteCompressedMessages(codec CompressionCodec, msgs ...Message) (nbytes int, err error) {
  991. nbytes, _, _, _, err = c.writeCompressedMessages(codec, msgs...)
  992. return
  993. }
  994. // WriteCompressedMessagesAt writes a batch of messages to the connection's topic
  995. // and partition, returning the number of bytes written, partition and offset numbers
  996. // and timestamp assigned by the kafka broker to the message set. The write is an atomic
  997. // operation, it either fully succeeds or fails.
  998. //
  999. // If the compression codec is not nil, the messages will be compressed.
  1000. func (c *Conn) WriteCompressedMessagesAt(codec CompressionCodec, msgs ...Message) (nbytes int, partition int32, offset int64, appendTime time.Time, err error) {
  1001. return c.writeCompressedMessages(codec, msgs...)
  1002. }
  1003. func (c *Conn) writeCompressedMessages(codec CompressionCodec, msgs ...Message) (nbytes int, partition int32, offset int64, appendTime time.Time, err error) {
  1004. if len(msgs) == 0 {
  1005. return
  1006. }
  1007. writeTime := time.Now()
  1008. for i, msg := range msgs {
  1009. // users may believe they can set the Topic and/or Partition
  1010. // on the kafka message.
  1011. if msg.Topic != "" && msg.Topic != c.topic {
  1012. err = errInvalidWriteTopic
  1013. return
  1014. }
  1015. if msg.Partition != 0 {
  1016. err = errInvalidWritePartition
  1017. return
  1018. }
  1019. if msg.Time.IsZero() {
  1020. msgs[i].Time = writeTime
  1021. }
  1022. nbytes += len(msg.Key) + len(msg.Value)
  1023. }
  1024. var produceVersion apiVersion
  1025. if produceVersion, err = c.negotiateVersion(produce, v2, v3, v7); err != nil {
  1026. return
  1027. }
  1028. err = c.writeOperation(
  1029. func(deadline time.Time, id int32) error {
  1030. now := time.Now()
  1031. deadline = adjustDeadlineForRTT(deadline, now, defaultRTT)
  1032. switch produceVersion {
  1033. case v7:
  1034. recordBatch, err :=
  1035. newRecordBatch(
  1036. codec,
  1037. msgs...,
  1038. )
  1039. if err != nil {
  1040. return err
  1041. }
  1042. return c.wb.writeProduceRequestV7(
  1043. id,
  1044. c.clientID,
  1045. c.topic,
  1046. c.partition,
  1047. deadlineToTimeout(deadline, now),
  1048. int16(atomic.LoadInt32(&c.requiredAcks)),
  1049. c.transactionalID,
  1050. recordBatch,
  1051. )
  1052. case v3:
  1053. recordBatch, err :=
  1054. newRecordBatch(
  1055. codec,
  1056. msgs...,
  1057. )
  1058. if err != nil {
  1059. return err
  1060. }
  1061. return c.wb.writeProduceRequestV3(
  1062. id,
  1063. c.clientID,
  1064. c.topic,
  1065. c.partition,
  1066. deadlineToTimeout(deadline, now),
  1067. int16(atomic.LoadInt32(&c.requiredAcks)),
  1068. c.transactionalID,
  1069. recordBatch,
  1070. )
  1071. default:
  1072. return c.wb.writeProduceRequestV2(
  1073. codec,
  1074. id,
  1075. c.clientID,
  1076. c.topic,
  1077. c.partition,
  1078. deadlineToTimeout(deadline, now),
  1079. int16(atomic.LoadInt32(&c.requiredAcks)),
  1080. msgs...,
  1081. )
  1082. }
  1083. },
  1084. func(deadline time.Time, size int) error {
  1085. return expectZeroSize(readArrayWith(&c.rbuf, size, func(r *bufio.Reader, size int) (int, error) {
  1086. // Skip the topic, we've produced the message to only one topic,
  1087. // no need to waste resources loading it in memory.
  1088. size, err := discardString(r, size)
  1089. if err != nil {
  1090. return size, err
  1091. }
  1092. // Read the list of partitions, there should be only one since
  1093. // we've produced a message to a single partition.
  1094. size, err = readArrayWith(r, size, func(r *bufio.Reader, size int) (int, error) {
  1095. switch produceVersion {
  1096. case v7:
  1097. var p produceResponsePartitionV7
  1098. size, err := p.readFrom(r, size)
  1099. if err == nil && p.ErrorCode != 0 {
  1100. err = Error(p.ErrorCode)
  1101. }
  1102. if err == nil {
  1103. partition = p.Partition
  1104. offset = p.Offset
  1105. appendTime = time.Unix(0, p.Timestamp*int64(time.Millisecond))
  1106. }
  1107. return size, err
  1108. default:
  1109. var p produceResponsePartitionV2
  1110. size, err := p.readFrom(r, size)
  1111. if err == nil && p.ErrorCode != 0 {
  1112. err = Error(p.ErrorCode)
  1113. }
  1114. if err == nil {
  1115. partition = p.Partition
  1116. offset = p.Offset
  1117. appendTime = time.Unix(0, p.Timestamp*int64(time.Millisecond))
  1118. }
  1119. return size, err
  1120. }
  1121. })
  1122. if err != nil {
  1123. return size, err
  1124. }
  1125. // The response is trailed by the throttle time, also skipping
  1126. // since it's not interesting here.
  1127. return discardInt32(r, size)
  1128. }))
  1129. },
  1130. )
  1131. if err != nil {
  1132. nbytes = 0
  1133. }
  1134. return
  1135. }
  1136. // SetRequiredAcks sets the number of acknowledges from replicas that the
  1137. // connection requests when producing messages.
  1138. func (c *Conn) SetRequiredAcks(n int) error {
  1139. switch n {
  1140. case -1, 1:
  1141. atomic.StoreInt32(&c.requiredAcks, int32(n))
  1142. return nil
  1143. default:
  1144. return InvalidRequiredAcks
  1145. }
  1146. }
  1147. func (c *Conn) writeRequest(apiKey apiKey, apiVersion apiVersion, correlationID int32, req request) error {
  1148. hdr := c.requestHeader(apiKey, apiVersion, correlationID)
  1149. hdr.Size = (hdr.size() + req.size()) - 4
  1150. hdr.writeTo(&c.wb)
  1151. req.writeTo(&c.wb)
  1152. return c.wbuf.Flush()
  1153. }
  1154. func (c *Conn) readResponse(size int, res interface{}) error {
  1155. size, err := read(&c.rbuf, size, res)
  1156. if err != nil {
  1157. var kafkaError Error
  1158. if errors.As(err, &kafkaError) {
  1159. size, err = discardN(&c.rbuf, size, size)
  1160. }
  1161. }
  1162. return expectZeroSize(size, err)
  1163. }
  1164. func (c *Conn) peekResponseSizeAndID() (int32, int32, error) {
  1165. b, err := c.rbuf.Peek(8)
  1166. if err != nil {
  1167. return 0, 0, err
  1168. }
  1169. size, id := makeInt32(b[:4]), makeInt32(b[4:])
  1170. return size, id, nil
  1171. }
  1172. func (c *Conn) skipResponseSizeAndID() {
  1173. c.rbuf.Discard(8)
  1174. }
  1175. func (c *Conn) readDeadline() time.Time {
  1176. return c.rdeadline.deadline()
  1177. }
  1178. func (c *Conn) writeDeadline() time.Time {
  1179. return c.wdeadline.deadline()
  1180. }
  1181. func (c *Conn) readOperation(write func(time.Time, int32) error, read func(time.Time, int) error) error {
  1182. return c.do(&c.rdeadline, write, read)
  1183. }
  1184. func (c *Conn) writeOperation(write func(time.Time, int32) error, read func(time.Time, int) error) error {
  1185. return c.do(&c.wdeadline, write, read)
  1186. }
  1187. func (c *Conn) enter() {
  1188. atomic.AddInt32(&c.inflight, +1)
  1189. }
  1190. func (c *Conn) leave() {
  1191. atomic.AddInt32(&c.inflight, -1)
  1192. }
  1193. func (c *Conn) concurrency() int {
  1194. return int(atomic.LoadInt32(&c.inflight))
  1195. }
  1196. func (c *Conn) do(d *connDeadline, write func(time.Time, int32) error, read func(time.Time, int) error) error {
  1197. id, err := c.doRequest(d, write)
  1198. if err != nil {
  1199. return err
  1200. }
  1201. deadline, size, lock, err := c.waitResponse(d, id)
  1202. if err != nil {
  1203. return err
  1204. }
  1205. if err = read(deadline, size); err != nil {
  1206. var kafkaError Error
  1207. if !errors.As(err, &kafkaError) {
  1208. c.conn.Close()
  1209. }
  1210. }
  1211. d.unsetConnReadDeadline()
  1212. lock.Unlock()
  1213. return err
  1214. }
  1215. func (c *Conn) doRequest(d *connDeadline, write func(time.Time, int32) error) (id int32, err error) {
  1216. c.enter()
  1217. c.wlock.Lock()
  1218. c.correlationID++
  1219. id = c.correlationID
  1220. err = write(d.setConnWriteDeadline(c.conn), id)
  1221. d.unsetConnWriteDeadline()
  1222. if err != nil {
  1223. // When an error occurs there's no way to know if the connection is in a
  1224. // recoverable state so we're better off just giving up at this point to
  1225. // avoid any risk of corrupting the following operations.
  1226. c.conn.Close()
  1227. c.leave()
  1228. }
  1229. c.wlock.Unlock()
  1230. return
  1231. }
  1232. func (c *Conn) waitResponse(d *connDeadline, id int32) (deadline time.Time, size int, lock *sync.Mutex, err error) {
  1233. for {
  1234. var rsz int32
  1235. var rid int32
  1236. c.rlock.Lock()
  1237. deadline = d.setConnReadDeadline(c.conn)
  1238. rsz, rid, err = c.peekResponseSizeAndID()
  1239. if err != nil {
  1240. d.unsetConnReadDeadline()
  1241. c.conn.Close()
  1242. c.rlock.Unlock()
  1243. break
  1244. }
  1245. if id == rid {
  1246. c.skipResponseSizeAndID()
  1247. size, lock = int(rsz-4), &c.rlock
  1248. // Don't unlock the read mutex to yield ownership to the caller.
  1249. break
  1250. }
  1251. if c.concurrency() == 1 {
  1252. // If the goroutine is the only one waiting on this connection it
  1253. // should be impossible to read a correlation id different from the
  1254. // one it expects. This is a sign that the data we are reading on
  1255. // the wire is corrupted and the connection needs to be closed.
  1256. err = io.ErrNoProgress
  1257. c.rlock.Unlock()
  1258. break
  1259. }
  1260. // Optimistically release the read lock if a response has already
  1261. // been received but the current operation is not the target for it.
  1262. c.rlock.Unlock()
  1263. }
  1264. c.leave()
  1265. return
  1266. }
  1267. func (c *Conn) requestHeader(apiKey apiKey, apiVersion apiVersion, correlationID int32) requestHeader {
  1268. return requestHeader{
  1269. ApiKey: int16(apiKey),
  1270. ApiVersion: int16(apiVersion),
  1271. CorrelationID: correlationID,
  1272. ClientID: c.clientID,
  1273. }
  1274. }
  1275. func (c *Conn) ApiVersions() ([]ApiVersion, error) {
  1276. deadline := &c.rdeadline
  1277. if deadline.deadline().IsZero() {
  1278. // ApiVersions is called automatically when API version negotiation
  1279. // needs to happen, so we are not guaranteed that a read deadline has
  1280. // been set yet. Fallback to use the write deadline in case it was
  1281. // set, for example when version negotiation is initiated during a
  1282. // produce request.
  1283. deadline = &c.wdeadline
  1284. }
  1285. id, err := c.doRequest(deadline, func(_ time.Time, id int32) error {
  1286. h := requestHeader{
  1287. ApiKey: int16(apiVersions),
  1288. ApiVersion: int16(v0),
  1289. CorrelationID: id,
  1290. ClientID: c.clientID,
  1291. }
  1292. h.Size = (h.size() - 4)
  1293. h.writeTo(&c.wb)
  1294. return c.wbuf.Flush()
  1295. })
  1296. if err != nil {
  1297. return nil, err
  1298. }
  1299. _, size, lock, err := c.waitResponse(deadline, id)
  1300. if err != nil {
  1301. return nil, err
  1302. }
  1303. defer lock.Unlock()
  1304. var errorCode int16
  1305. if size, err = readInt16(&c.rbuf, size, &errorCode); err != nil {
  1306. return nil, err
  1307. }
  1308. var arrSize int32
  1309. if size, err = readInt32(&c.rbuf, size, &arrSize); err != nil {
  1310. return nil, err
  1311. }
  1312. r := make([]ApiVersion, arrSize)
  1313. for i := 0; i < int(arrSize); i++ {
  1314. if size, err = readInt16(&c.rbuf, size, &r[i].ApiKey); err != nil {
  1315. return nil, err
  1316. }
  1317. if size, err = readInt16(&c.rbuf, size, &r[i].MinVersion); err != nil {
  1318. return nil, err
  1319. }
  1320. if size, err = readInt16(&c.rbuf, size, &r[i].MaxVersion); err != nil {
  1321. return nil, err
  1322. }
  1323. }
  1324. if errorCode != 0 {
  1325. return r, Error(errorCode)
  1326. }
  1327. return r, nil
  1328. }
  1329. // connDeadline is a helper type to implement read/write deadline management on
  1330. // the kafka connection.
  1331. type connDeadline struct {
  1332. mutex sync.Mutex
  1333. value time.Time
  1334. rconn net.Conn
  1335. wconn net.Conn
  1336. }
  1337. func (d *connDeadline) deadline() time.Time {
  1338. d.mutex.Lock()
  1339. t := d.value
  1340. d.mutex.Unlock()
  1341. return t
  1342. }
  1343. func (d *connDeadline) setDeadline(t time.Time) {
  1344. d.mutex.Lock()
  1345. d.value = t
  1346. if d.rconn != nil {
  1347. d.rconn.SetReadDeadline(t)
  1348. }
  1349. if d.wconn != nil {
  1350. d.wconn.SetWriteDeadline(t)
  1351. }
  1352. d.mutex.Unlock()
  1353. }
  1354. func (d *connDeadline) setConnReadDeadline(conn net.Conn) time.Time {
  1355. d.mutex.Lock()
  1356. deadline := d.value
  1357. d.rconn = conn
  1358. d.rconn.SetReadDeadline(deadline)
  1359. d.mutex.Unlock()
  1360. return deadline
  1361. }
  1362. func (d *connDeadline) setConnWriteDeadline(conn net.Conn) time.Time {
  1363. d.mutex.Lock()
  1364. deadline := d.value
  1365. d.wconn = conn
  1366. d.wconn.SetWriteDeadline(deadline)
  1367. d.mutex.Unlock()
  1368. return deadline
  1369. }
  1370. func (d *connDeadline) unsetConnReadDeadline() {
  1371. d.mutex.Lock()
  1372. d.rconn = nil
  1373. d.mutex.Unlock()
  1374. }
  1375. func (d *connDeadline) unsetConnWriteDeadline() {
  1376. d.mutex.Lock()
  1377. d.wconn = nil
  1378. d.mutex.Unlock()
  1379. }
  1380. // saslHandshake sends the SASL handshake message. This will determine whether
  1381. // the Mechanism is supported by the cluster. If it's not, this function will
  1382. // error out with UnsupportedSASLMechanism.
  1383. //
  1384. // If the mechanism is unsupported, the handshake request will reply with the
  1385. // list of the cluster's configured mechanisms, which could potentially be used
  1386. // to facilitate negotiation. At the moment, we are not negotiating the
  1387. // mechanism as we believe that brokers are usually known to the client, and
  1388. // therefore the client should already know which mechanisms are supported.
  1389. //
  1390. // See http://kafka.apache.org/protocol.html#The_Messages_SaslHandshake
  1391. func (c *Conn) saslHandshake(mechanism string) error {
  1392. // The wire format for V0 and V1 is identical, but the version
  1393. // number will affect how the SASL authentication
  1394. // challenge/responses are sent
  1395. var resp saslHandshakeResponseV0
  1396. version, err := c.negotiateVersion(saslHandshake, v0, v1)
  1397. if err != nil {
  1398. return err
  1399. }
  1400. err = c.writeOperation(
  1401. func(deadline time.Time, id int32) error {
  1402. return c.writeRequest(saslHandshake, version, id, &saslHandshakeRequestV0{Mechanism: mechanism})
  1403. },
  1404. func(deadline time.Time, size int) error {
  1405. return expectZeroSize(func() (int, error) {
  1406. return (&resp).readFrom(&c.rbuf, size)
  1407. }())
  1408. },
  1409. )
  1410. if err == nil && resp.ErrorCode != 0 {
  1411. err = Error(resp.ErrorCode)
  1412. }
  1413. return err
  1414. }
  1415. // saslAuthenticate sends the SASL authenticate message. This function must
  1416. // be immediately preceded by a successful saslHandshake.
  1417. //
  1418. // See http://kafka.apache.org/protocol.html#The_Messages_SaslAuthenticate
  1419. func (c *Conn) saslAuthenticate(data []byte) ([]byte, error) {
  1420. // if we sent a v1 handshake, then we must encapsulate the authentication
  1421. // request in a saslAuthenticateRequest. otherwise, we read and write raw
  1422. // bytes.
  1423. version, err := c.negotiateVersion(saslHandshake, v0, v1)
  1424. if err != nil {
  1425. return nil, err
  1426. }
  1427. if version == v1 {
  1428. var request = saslAuthenticateRequestV0{Data: data}
  1429. var response saslAuthenticateResponseV0
  1430. err := c.writeOperation(
  1431. func(deadline time.Time, id int32) error {
  1432. return c.writeRequest(saslAuthenticate, v0, id, request)
  1433. },
  1434. func(deadline time.Time, size int) error {
  1435. return expectZeroSize(func() (remain int, err error) {
  1436. return (&response).readFrom(&c.rbuf, size)
  1437. }())
  1438. },
  1439. )
  1440. if err == nil && response.ErrorCode != 0 {
  1441. err = Error(response.ErrorCode)
  1442. }
  1443. return response.Data, err
  1444. }
  1445. // fall back to opaque bytes on the wire. the broker is expecting these if
  1446. // it just processed a v0 sasl handshake.
  1447. c.wb.writeInt32(int32(len(data)))
  1448. if _, err := c.wb.Write(data); err != nil {
  1449. return nil, err
  1450. }
  1451. if err := c.wb.Flush(); err != nil {
  1452. return nil, err
  1453. }
  1454. var respLen int32
  1455. if _, err := readInt32(&c.rbuf, 4, &respLen); err != nil {
  1456. return nil, err
  1457. }
  1458. resp, _, err := readNewBytes(&c.rbuf, int(respLen), int(respLen))
  1459. return resp, err
  1460. }