Merge pull request #1 from gostudentorg/DT-14757

DT-14757: Fix: writers closing issue

writer.go

@@ -663,7 +663,10 @@ func (w *Writer) WriteMessages(ctx context.Context, msgs ...Message) error {
  assignments[key] = append(assignments[key], int32(i))
  }
 
- batches := w.batchMessages(msgs, assignments)
+ batches, err := w.batchMessages(msgs, assignments)
+ if err != nil {
+ return err
+ }
  if w.Async {
  return nil
  }
@@ -695,7 +698,7 @@ func (w *Writer) WriteMessages(ctx context.Context, msgs ...Message) error {
  return werr
 }
 
-func (w *Writer) batchMessages(messages []Message, assignments map[topicPartition][]int32) map[*writeBatch][]int32 {
+func (w *Writer) batchMessages(messages []Message, assignments map[topicPartition][]int32) (map[*writeBatch][]int32, error) {
  var batches map[*writeBatch][]int32
  if !w.Async {
  batches = make(map[*writeBatch][]int32, len(assignments))
@@ -704,6 +707,10 @@ func (w *Writer) batchMessages(messages []Message, assignments map[topicPartitio
  w.mutex.Lock()
  defer w.mutex.Unlock()
 
+ if w.closed {
+ return nil, io.ErrClosedPipe
+ }
+
  if w.writers == nil {
  w.writers = map[topicPartition]*partitionWriter{}
  }
@@ -721,7 +728,7 @@ func (w *Writer) batchMessages(messages []Message, assignments map[topicPartitio
  }
  }
 
- return batches
+ return batches, nil
 }
 
 func (w *Writer) produce(key topicPartition, batch *writeBatch) (*ProduceResponse, error) {

writer_test.go

@@ -191,6 +191,10 @@ func TestWriter(t *testing.T) {
  scenario: "test write message with writer data",
  function: testWriteMessageWithWriterData,
  },
+ {
+ scenario: "test no new partition writers after close",
+ function: testWriterNoNewPartitionWritersAfterClose,
+ },
  }
 
  for _, test := range tests {
@@ -1030,6 +1034,46 @@ func testWriterOverrideConfigStats(t *testing.T) {
  }
 }
 
+func testWriterNoNewPartitionWritersAfterClose(t *testing.T) {
+ ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+ defer cancel()
+ topic1 := makeTopic()
+ createTopic(t, topic1, 1)
+ defer deleteTopic(t, topic1)
+
+ w := newTestWriter(WriterConfig{
+ Topic: topic1,
+ })
+ defer w.Close() // try and close anyway after test finished
+
+ // using balancer to close writer right between first mutex is released and second mutex is taken to make map of partition writers
+ w.Balancer = mockBalancerFunc(func(m Message, i ...int) int {
+ go w.Close() // close is blocking so run in goroutine
+ for { // wait until writer is marked as closed
+ w.mutex.Lock()
+ if w.closed {
+ w.mutex.Unlock()
+ break
+ }
+ w.mutex.Unlock()
+ }
+ return 0
+ })
+
+ msg := Message{Value: []byte("Hello World")} // no topic
+
+ if err := w.WriteMessages(ctx, msg); !errors.Is(err, io.ErrClosedPipe) {
+ t.Errorf("expected error: %v got: %v", io.ErrClosedPipe, err)
+ return
+ }
+}
+
+type mockBalancerFunc func(msg Message, partitions ...int) (partition int)
+
+func (b mockBalancerFunc) Balance(msg Message, partitions ...int) int {
+ return b(msg, partitions...)
+}
+
 type staticBalancer struct {
  partition int
 }