github.com/segmentio/kafka-go@v0.4.48-0.20240318174348-3f6244eb34fd/protocol/conn.go (about)

     1  package protocol
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"net"
     7  	"sync/atomic"
     8  	"time"
     9  )
    10  
    11  type Conn struct {
    12  	buffer   *bufio.Reader
    13  	conn     net.Conn
    14  	clientID string
    15  	idgen    int32
    16  	versions atomic.Value // map[ApiKey]int16
    17  }
    18  
    19  func NewConn(conn net.Conn, clientID string) *Conn {
    20  	return &Conn{
    21  		buffer:   bufio.NewReader(conn),
    22  		conn:     conn,
    23  		clientID: clientID,
    24  	}
    25  }
    26  
    27  func (c *Conn) String() string {
    28  	return fmt.Sprintf("kafka://%s@%s->%s", c.clientID, c.LocalAddr(), c.RemoteAddr())
    29  }
    30  
    31  func (c *Conn) Close() error {
    32  	return c.conn.Close()
    33  }
    34  
    35  func (c *Conn) Discard(n int) (int, error) {
    36  	return c.buffer.Discard(n)
    37  }
    38  
    39  func (c *Conn) Peek(n int) ([]byte, error) {
    40  	return c.buffer.Peek(n)
    41  }
    42  
    43  func (c *Conn) Read(b []byte) (int, error) {
    44  	return c.buffer.Read(b)
    45  }
    46  
    47  func (c *Conn) Write(b []byte) (int, error) {
    48  	return c.conn.Write(b)
    49  }
    50  
    51  func (c *Conn) LocalAddr() net.Addr {
    52  	return c.conn.LocalAddr()
    53  }
    54  
    55  func (c *Conn) RemoteAddr() net.Addr {
    56  	return c.conn.RemoteAddr()
    57  }
    58  
    59  func (c *Conn) SetDeadline(t time.Time) error {
    60  	return c.conn.SetDeadline(t)
    61  }
    62  
    63  func (c *Conn) SetReadDeadline(t time.Time) error {
    64  	return c.conn.SetReadDeadline(t)
    65  }
    66  
    67  func (c *Conn) SetWriteDeadline(t time.Time) error {
    68  	return c.conn.SetWriteDeadline(t)
    69  }
    70  
    71  func (c *Conn) SetVersions(versions map[ApiKey]int16) {
    72  	connVersions := make(map[ApiKey]int16, len(versions))
    73  
    74  	for k, v := range versions {
    75  		connVersions[k] = v
    76  	}
    77  
    78  	c.versions.Store(connVersions)
    79  }
    80  
    81  func (c *Conn) RoundTrip(msg Message) (Message, error) {
    82  	correlationID := atomic.AddInt32(&c.idgen, +1)
    83  	versions, _ := c.versions.Load().(map[ApiKey]int16)
    84  	apiVersion := versions[msg.ApiKey()]
    85  
    86  	if p, _ := msg.(PreparedMessage); p != nil {
    87  		p.Prepare(apiVersion)
    88  	}
    89  
    90  	if raw, ok := msg.(RawExchanger); ok && raw.Required(versions) {
    91  		return raw.RawExchange(c)
    92  	}
    93  
    94  	return RoundTrip(c, apiVersion, correlationID, c.clientID, msg)
    95  }
    96  
    97  var (
    98  	_ net.Conn       = (*Conn)(nil)
    99  	_ bufferedReader = (*Conn)(nil)
   100  )