stop sending if the connenction is closed

Author: gunli

pulsar/consumer_multitopic_test.go

@@ -332,7 +332,8 @@ func (dummyConnection) SendRequestNoWait(_ *pb.BaseCommand) error {
 	return nil
 }
 
-func (dummyConnection) WriteData(_ context.Context, _ internal.Buffer) {
+func (dummyConnection) WriteData(_ context.Context, _ internal.Buffer) error {
+	return nil
 }
 
 func (dummyConnection) RegisterListener(_ uint64, _ internal.ConnectionListener) error {

pulsar/internal/connection.go

@@ -79,7 +79,7 @@ type ConnectionListener interface {
 type Connection interface {
 	SendRequest(requestID uint64, req *pb.BaseCommand, callback func(*pb.BaseCommand, error))
 	SendRequestNoWait(req *pb.BaseCommand) error
-	WriteData(ctx context.Context, data Buffer)
+	WriteData(ctx context.Context, data Buffer) error
 	RegisterListener(id uint64, listener ConnectionListener) error
 	UnregisterListener(id uint64)
 	AddConsumeHandler(id uint64, handler ConsumerHandler) error
@@ -456,26 +456,38 @@ func (c *connection) runPingCheck(pingCheckTicker *time.Ticker) {
 	}
 }
 
-func (c *connection) WriteData(ctx context.Context, data Buffer) {
+func (c *connection) WriteData(ctx context.Context, data Buffer) error {
+	// If the connection is closed, we should not write on it
+	if c.getState() != connectionReady {
+		c.log.Debug("Connection was already closed")
+		return errConnectionClosed
+	}
+
 	select {
 	case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 		// Channel is not full
-		return
+		return nil
 	case <-ctx.Done():
 		c.log.Debug("Write data context cancelled")
-		return
+		return ctx.Err()
 	default:
 		// Channel full, fallback to probe if connection is closed
 	}
 
+	// check if the connection is closed again
+	if c.getState() != connectionReady {
+		c.log.Debug("Connection was already closed")
+		return errConnectionClosed
+	}
+
 	for {
 		select {
 		case c.writeRequestsCh <- &dataRequest{ctx: ctx, data: data}:
 			// Successfully wrote on the channel
-			return
+			return nil
 		case <-ctx.Done():
 			c.log.Debug("Write data context cancelled")
-			return
+			return ctx.Err()
 		case <-time.After(100 * time.Millisecond):
 			// The channel is either:
 			// 1. blocked, in which case we need to wait until we have space
@@ -484,7 +496,7 @@ func (c *connection) WriteData(ctx context.Context, data Buffer) {
 
 			if c.getState() != connectionReady {
 				c.log.Debug("Connection was already closed")
-				return
+				return errConnectionClosed
 			}
 		}
 	}

pulsar/internal/connection_reader.go

@@ -21,8 +21,9 @@ import (
 	"fmt"
 	"io"
 
-	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
 	"google.golang.org/protobuf/proto"
+
+	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
 )
 
 type connectionReader struct {

pulsar/producer_partition.go

@@ -394,7 +394,10 @@ func (p *partitionProducer) grabCnx(assignedBrokerURL string) error {
 			pi.sentAt = time.Now()
 			pi.Unlock()
 			p.pendingQueue.Put(pi)
-			p._getConn().WriteData(pi.ctx, pi.buffer)
+			err = p._getConn().WriteData(pi.ctx, pi.buffer)
+			if err != nil {
+				p.log.WithError(err).Warn("failed to write data, it will be retried later")
+			}
 
 			if pi == lastViewItem {
 				break
@@ -908,7 +911,13 @@ func (p *partitionProducer) writeData(buffer internal.Buffer, sequenceID uint64,
 			sequenceID:   sequenceID,
 			sendRequests: callbacks,
 		})
-		p._getConn().WriteData(ctx, buffer)
+
+		// If the connection is closed, WriteData() will failed, but it is fine, the buffer is still kept in p.pendingQueue,
+		// it will be sent out or timeout finally.
+		err := p._getConn().WriteData(ctx, buffer)
+		if err != nil {
+			p.log.WithError(err).Warn("failed to write data, it will be retried later")
+		}
 	}
 }