Use ping-pong buffer for batch conn

To prevent writes from getting blocked while flushing the batch,
switching to a ping-pong buffer.

Add some more test to explicitly force write batch interval and write
batch size flush.

Author: boks1971

udp/batchconn.go

@@ -32,35 +32,248 @@ type BatchPacketConn interface {
 	io.Closer
 }
 
+// ---------------------------------
+
+type messageBatch struct {
+	size      int
+	batchConn BatchPacketConn
+
+	mu       sync.Mutex
+	messages []ipv4.Message
+	writePos int
+}
+
+func newMessageBatch(size int, batchConn BatchPacketConn) *messageBatch {
+	m := &messageBatch{
+		size:      size,
+		batchConn: batchConn,
+	}
+	m.init()
+
+	return m
+}
+
+func (m *messageBatch) init() {
+	m.messages = make([]ipv4.Message, m.size)
+	for i := range m.messages {
+		m.messages[i].Buffers = [][]byte{make([]byte, sendMTU)}
+	}
+}
+
+func (m *messageBatch) isFull() bool {
+	return m.writePos == m.size
+}
+
+func (m *messageBatch) EnqueueMessage(buf []byte, raddr net.Addr) (int, bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if len(buf) == 0 || m.isFull() {
+		return 0, m.isFull()
+	}
+
+	msg := &m.messages[m.writePos]
+	// reset buffers
+	msg.Buffers = msg.Buffers[:1]
+	msg.Buffers[0] = msg.Buffers[0][:cap(msg.Buffers[0])]
+
+	if raddr != nil {
+		msg.Addr = raddr
+	}
+	if n := copy(msg.Buffers[0], buf); n < len(buf) {
+		extraBuffer := make([]byte, len(buf)-n)
+		copy(extraBuffer, buf[n:])
+		msg.Buffers = append(msg.Buffers, extraBuffer)
+	} else {
+		msg.Buffers[0] = msg.Buffers[0][:n]
+	}
+	m.writePos++
+
+	return len(buf), m.isFull()
+}
+
+func (m *messageBatch) Flush() {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	var txN int
+	for txN < m.writePos {
+		n, err := m.batchConn.WriteBatch(m.messages[txN:m.writePos], 0)
+		if err != nil {
+			break
+		}
+		txN += n
+	}
+
+	m.writePos = 0
+}
+
+// ---------------------------------
+
+type pingPong struct {
+	mu            sync.Mutex
+	batches       [2]*messageBatch
+	writeBatchIdx int
+	readBatchIdx  int
+	flushPending  bool
+
+	writeReady     chan struct{}
+	flushCycleDone chan struct{}
+	flusherDone    chan struct{}
+
+	closed int32
+}
+
+func newPingPong(size int, interval time.Duration, batchConn BatchPacketConn) *pingPong {
+	p := &pingPong{
+		writeReady:     make(chan struct{}),
+		flushCycleDone: make(chan struct{}),
+		flusherDone:    make(chan struct{}),
+	}
+	for i := 0; i < len(p.batches); i++ {
+		p.batches[i] = newMessageBatch(size, batchConn)
+	}
+
+	go p.flusher(interval)
+
+	return p
+}
+
+func (p *pingPong) Close() {
+	atomic.StoreInt32(&p.closed, 1)
+
+	select {
+	case p.writeReady <- struct{}{}:
+	default:
+	}
+
+	<-p.flusherDone
+}
+
+func (p *pingPong) EnqueueMessage(buf []byte, raddr net.Addr) int {
+	p.mu.Lock()
+	var (
+		writeBatch *messageBatch
+		n          int
+		isFull     bool
+	)
+	for {
+		if writeBatch = p.getWriteBatch(); writeBatch != nil {
+			n, isFull = writeBatch.EnqueueMessage(buf, raddr)
+			if n == len(buf) {
+				break
+			}
+		}
+
+		p.mu.Unlock()
+		select {
+		case <-p.flushCycleDone:
+		case <-time.After(100 * time.Microsecond):
+		}
+
+		if atomic.LoadInt32(&p.closed) == 1 {
+			return 0
+		}
+
+		p.mu.Lock()
+	}
+	p.mu.Unlock()
+
+	// enqueuing given message fills up the write batch, queue up a flush
+	if isFull {
+		select {
+		case p.writeReady <- struct{}{}:
+		default:
+		}
+	}
+
+	return n
+}
+
+func (p *pingPong) getWriteBatch() *messageBatch {
+	if p.writeBatchIdx == p.readBatchIdx && p.flushPending {
+		return nil
+	}
+
+	return p.batches[p.writeBatchIdx]
+}
+
+func (p *pingPong) updateWriteBatchAndGetReadBatch() *messageBatch {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if p.writeBatchIdx != p.readBatchIdx || p.flushPending {
+		return nil
+	}
+
+	p.writeBatchIdx ^= 1
+	p.flushPending = true
+
+	return p.batches[p.readBatchIdx]
+}
+
+func (p *pingPong) updateReadBatch() {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	if !p.flushPending {
+		return
+	}
+
+	p.readBatchIdx ^= 1
+	p.flushPending = false
+
+	select {
+	case p.flushCycleDone <- struct{}{}:
+	default:
+	}
+}
+
+func (p *pingPong) flusher(interval time.Duration) {
+	defer close(p.flusherDone)
+
+	writeTicker := time.NewTicker(interval / 2)
+	defer writeTicker.Stop()
+
+	lastFlushAt := time.Now().Add(-interval)
+	for atomic.LoadInt32(&p.closed) != 1 {
+		select {
+		case <-writeTicker.C:
+			if time.Since(lastFlushAt) < interval {
+				continue
+			}
+		case <-p.writeReady:
+		}
+
+		readBatch := p.updateWriteBatchAndGetReadBatch()
+		if readBatch == nil {
+			continue
+		}
+
+		readBatch.Flush()
+		p.updateReadBatch()
+
+		lastFlushAt = time.Now()
+	}
+}
+
+// ------------------------------
+
 // BatchConn uses ipv4/v6.NewPacketConn to wrap a net.PacketConn to write/read messages in batch,
 // only available in linux. In other platform, it will use single Write/Read as same as net.Conn.
 type BatchConn struct {
 	net.PacketConn
 
 	batchConn BatchPacketConn
 
-	batchWriteMutex    sync.Mutex
-	batchWriteMessages []ipv4.Message
-	batchWritePos      int
-	batchWriteLast     time.Time
-
-	batchWriteSize     int
-	batchWriteInterval time.Duration
-
-	closed int32
+	// ping-pong the batches to be able to accept new packets while a batch is written to socket
+	batchPingPong *pingPong
 }
 
 // NewBatchConn creates a *BatchConn from net.PacketConn with batch configs.
 func NewBatchConn(conn net.PacketConn, batchWriteSize int, batchWriteInterval time.Duration) *BatchConn {
 	bc := &BatchConn{
-		PacketConn:         conn,
-		batchWriteLast:     time.Now(),
-		batchWriteInterval: batchWriteInterval,
-		batchWriteSize:     batchWriteSize,
-		batchWriteMessages: make([]ipv4.Message, batchWriteSize),
-	}
-	for i := range bc.batchWriteMessages {
-		bc.batchWriteMessages[i].Buffers = [][]byte{make([]byte, sendMTU)}
+		PacketConn: conn,
 	}
 
 	// batch write only supports linux
@@ -70,34 +283,19 @@ func NewBatchConn(conn net.PacketConn, batchWriteSize int, batchWriteInterval ti
 		} else if pc6 := ipv6.NewPacketConn(conn); pc6 != nil {
 			bc.batchConn = pc6
 		}
-	}
 
-	if bc.batchConn != nil {
-		go func() {
-			writeTicker := time.NewTicker(batchWriteInterval / 2)
-			defer writeTicker.Stop()
-			for atomic.LoadInt32(&bc.closed) != 1 {
-				<-writeTicker.C
-				bc.batchWriteMutex.Lock()
-				if bc.batchWritePos > 0 && time.Since(bc.batchWriteLast) >= bc.batchWriteInterval {
-					_ = bc.flush()
-				}
-				bc.batchWriteMutex.Unlock()
-			}
-		}()
+		bc.batchPingPong = newPingPong(batchWriteSize, batchWriteInterval, bc.batchConn)
 	}
 
 	return bc
 }
 
 // Close batchConn and the underlying PacketConn.
 func (c *BatchConn) Close() error {
-	atomic.StoreInt32(&c.closed, 1)
-	c.batchWriteMutex.Lock()
-	if c.batchWritePos > 0 {
-		_ = c.flush()
+	if c.batchPingPong != nil {
+		c.batchPingPong.Close()
 	}
-	c.batchWriteMutex.Unlock()
+
 	if c.batchConn != nil {
 		return c.batchConn.Close()
 	}
@@ -111,35 +309,7 @@ func (c *BatchConn) WriteTo(b []byte, addr net.Addr) (int, error) {
 		return c.PacketConn.WriteTo(b, addr)
 	}
 
-	return c.enqueueMessage(b, addr)
-}
-
-func (c *BatchConn) enqueueMessage(buf []byte, raddr net.Addr) (int, error) {
-	var err error
-	c.batchWriteMutex.Lock()
-	defer c.batchWriteMutex.Unlock()
-
-	msg := &c.batchWriteMessages[c.batchWritePos]
-	// reset buffers
-	msg.Buffers = msg.Buffers[:1]
-	msg.Buffers[0] = msg.Buffers[0][:cap(msg.Buffers[0])]
-
-	c.batchWritePos++
-	if raddr != nil {
-		msg.Addr = raddr
-	}
-	if n := copy(msg.Buffers[0], buf); n < len(buf) {
-		extraBuffer := make([]byte, len(buf)-n)
-		copy(extraBuffer, buf[n:])
-		msg.Buffers = append(msg.Buffers, extraBuffer)
-	} else {
-		msg.Buffers[0] = msg.Buffers[0][:n]
-	}
-	if c.batchWritePos == c.batchWriteSize {
-		err = c.flush()
-	}
-
-	return len(buf), err
+	return c.batchPingPong.EnqueueMessage(b, addr), nil
 }
 
 // ReadBatch reads messages in batch, return length of message readed and error.
@@ -158,21 +328,3 @@ func (c *BatchConn) ReadBatch(msgs []ipv4.Message, flags int) (int, error) {
 
 	return c.batchConn.ReadBatch(msgs, flags)
 }
-
-func (c *BatchConn) flush() error {
-	var writeErr error
-	var txN int
-	for txN < c.batchWritePos {
-		n, err := c.batchConn.WriteBatch(c.batchWriteMessages[txN:c.batchWritePos], 0)
-		if err != nil {
-			writeErr = err
-
-			break
-		}
-		txN += n
-	}
-	c.batchWritePos = 0
-	c.batchWriteLast = time.Now()
-
-	return writeErr
-}