stop writing if ctx is done

Author: gunli

pulsar/internal/connection.go

@@ -130,6 +130,11 @@ type request struct {
 	callback func(command *pb.BaseCommand, err error)
 }
 
+type dataRequest struct {
+	ctx  context.Context
+	data Buffer
+}
+
 type connection struct {
 	started           int32
 	connectionTimeout time.Duration
@@ -158,7 +163,7 @@ type connection struct {
 	incomingRequestsCh chan *request
 	closeCh            chan struct{}
 	readyCh            chan struct{}
-	writeRequestsCh    chan Buffer
+	writeRequestsCh    chan *dataRequest
 
 	pendingLock sync.Mutex
 	pendingReqs map[uint64]*request
@@ -210,7 +215,7 @@ func newConnection(opts connectionOptions) *connection {
 		// partition produces writing on a single connection. In general it's
 		// good to keep this above the number of partition producers assigned
 		// to a single connection.
-		writeRequestsCh:  make(chan Buffer, 256),
+		writeRequestsCh:  make(chan *dataRequest, 256),
 		listeners:        make(map[uint64]ConnectionListener),
 		consumerHandlers: make(map[uint64]ConsumerHandler),
 		metrics:          opts.metrics,
@@ -422,11 +427,11 @@ func (c *connection) run() {
 				return // TODO: this never gonna be happen
 			}
 			c.internalSendRequest(req)
-		case data := <-c.writeRequestsCh:
-			if data == nil {
+		case req := <-c.writeRequestsCh:
+			if req == nil {
 				return
 			}
-			c.internalWriteData(data)
+			c.internalWriteData(req.ctx, req.data)
 
 		case <-pingSendTicker.C:
 			c.sendPing()
@@ -453,7 +458,7 @@ func (c *connection) runPingCheck(pingCheckTicker *time.Ticker) {
 
 func (c *connection) WriteData(ctx context.Context, data Buffer) {
 	select {
-	case c.writeRequestsCh <- data:
+	case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 		// Channel is not full
 		return
 	case <-ctx.Done():
@@ -465,7 +470,7 @@ func (c *connection) WriteData(ctx context.Context, data Buffer) {
 
 	for {
 		select {
-		case c.writeRequestsCh <- data:
+		case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 			// Successfully wrote on the channel
 			return
 		case <-ctx.Done():
@@ -486,11 +491,25 @@ func (c *connection) WriteData(ctx context.Context, data Buffer) {
 
 }
 
-func (c *connection) internalWriteData(data Buffer) {
+func (c *connection) internalWriteData(ctx context.Context, data Buffer) {
 	c.log.Debug("Write data: ", data.ReadableBytes())
-	if _, err := c.cnx.Write(data.ReadableSlice()); err != nil {
-		c.log.WithError(err).Warn("Failed to write on connection")
-		c.Close()
+	if ctx == nil {
+		if _, err := c.cnx.Write(data.ReadableSlice()); err != nil {
+			c.log.WithError(err).Warn("Failed to write on connection")
+			c.Close()
+		}
+
+		return
+	}
+
+	select {
+	case <-ctx.Done():
+		return
+	default:
+		if _, err := c.cnx.Write(data.ReadableSlice()); err != nil {
+			c.log.WithError(err).Warn("Failed to write on connection")
+			c.Close()
+		}
 	}
 }
 
@@ -515,7 +534,7 @@ func (c *connection) writeCommand(cmd *pb.BaseCommand) {
 	}
 
 	c.writeBuffer.WrittenBytes(cmdSize)
-	c.internalWriteData(c.writeBuffer)
+	c.internalWriteData(nil, c.writeBuffer)
 }
 
 func (c *connection) receivedCommand(cmd *pb.BaseCommand, headersAndPayload Buffer) {