Optimize allocations for handshake and compaction

Author: noboruma

conn.go

@@ -573,7 +573,7 @@ func (c *Conn) Write(payload []byte) (int, error) {
 		return 0, err
 	}
 
-	ctx, cancel := c.contextWithClose(c.writeDeadline)
+	ctx, cancel := c.contextWithClose()
 	defer cancel()
 
 	err := c.writeApplicationData(ctx, []*dtlsflight.Outbound{
@@ -829,62 +829,94 @@ func (c *Conn) cacheHandshake(outbound *dtlsflight.Outbound, dtlsHandshake *hand
 	return nil
 }
 
-func (c *Conn) contextWithClose(ctx context.Context) (context.Context, context.CancelFunc) {
-	closeCtx, cancel := context.WithCancelCause(context.WithoutCancel(ctx))
-	go func() {
-		select {
-		case <-c.closed.Done():
-			cancel(context.Canceled)
-		case <-ctx.Done():
-			err := ctx.Err()
-			if err == nil {
-				err = context.DeadlineExceeded
-			}
-			cancel(err)
-		case <-closeCtx.Done():
-		}
-	}()
+func (c *Conn) contextWithClose() (context.Context, context.CancelFunc) {
+	ctx := context.Background()
+
+	var cancelDeadline context.CancelFunc = func() {}
+	if deadline, ok := c.readDeadline.Deadline(); ok {
+		ctx, cancelDeadline = context.WithDeadline(ctx, deadline)
+	}
+
+	closeCtx, cancelClose := context.WithCancelCause(ctx)
+	detachLifetime := context.AfterFunc(c.closed, func() {
+		cancelClose(c.closed.Err())
+	})
 
 	return closeCtx, func() {
-		cancel(context.Canceled)
+		detachLifetime()
+		cancelDeadline()
+		cancelClose(context.Canceled)
 	}
 }
 
 func (c *Conn) contextWithCloseAndWriteDeadline(ctx context.Context) (context.Context, context.CancelFunc) {
-	operationCtx, cancel := context.WithCancelCause(context.Background())
-	go func() {
-		select {
-		case <-c.closed.Done():
-			cancel(context.Canceled)
-		case <-c.writeDeadline.Done():
-			cancel(context.DeadlineExceeded)
-		case <-ctx.Done():
-			cancel(ctx.Err())
-		case <-operationCtx.Done():
+	if ctx == nil {
+		ctx = context.Background()
+	}
+
+	var cancelDeadline context.CancelFunc = func() {}
+	if deadline, ok := c.writeDeadline.Deadline(); ok {
+		ctx, cancelDeadline = context.WithDeadline(ctx, deadline)
+	}
+
+	operationCtx, cancelClose := context.WithCancelCause(ctx)
+	detachLifetime := context.AfterFunc(c.closed, func() {
+		err := c.closed.Err()
+		if err == nil {
+			err = context.Canceled
 		}
-	}()
+		cancelClose(err)
+	})
 
 	return operationCtx, func() {
-		cancel(context.Canceled)
+		detachLifetime()
+		cancelDeadline()
+		cancelClose(context.Canceled)
 	}
 }
 
 func (c *Conn) compactPreparedRecords(records []preparedRecord) []preparedDatagram {
-	datagrams := make([]preparedDatagram, 0, len(records))
-	current := preparedDatagram{}
+	if len(records) == 0 {
+		return []preparedDatagram{}
+	}
+
+	totalSize := 0
 	for _, record := range records {
-		if len(current.raw) > 0 && len(current.raw)+len(record.raw) >= c.maximumTransmissionUnit {
-			datagrams = append(datagrams, current)
-			current = preparedDatagram{}
+		totalSize += len(record.raw)
+	}
+
+	datagrams := make([]preparedDatagram, len(records))
+	flatRaw := make([]byte, 0, totalSize)
+
+	datagramIndex := 0
+	currentSize := 0
+	offset := 0
+
+	for _, record := range records {
+		recordSize := len(record.raw)
+
+		flatRaw = append(flatRaw, record.raw...)
+
+		if currentSize > 0 && currentSize+recordSize >= c.maximumTransmissionUnit {
+			datagrams[datagramIndex].raw = flatRaw[offset : offset+currentSize]
+			datagramIndex++
+			offset += currentSize
+			currentSize = 0
 		}
-		current.raw = append(current.raw, record.raw...)
+
+		currentSize += recordSize
+
 		if record.tracked != nil {
-			current.tracked = append(current.tracked, *record.tracked)
+			datagrams[datagramIndex].tracked = append(
+				datagrams[datagramIndex].tracked,
+				*record.tracked,
+			)
 		}
 	}
-	datagrams = append(datagrams, current)
 
-	return datagrams
+	datagrams[datagramIndex].raw = flatRaw[offset : offset+currentSize]
+
+	return datagrams[:datagramIndex+1]
 }
 
 func (c *Conn) prepareRecord(outbound *dtlsflight.Outbound) ([]byte, error) {
@@ -1214,42 +1246,76 @@ func selectHandshakeFragment(offsets map[uint32]uint32, raw []byte) (bool, error
 	return ok && length == header.FragmentLength, nil
 }
 
+var noContentFragments = [][]byte{ //nolint:gochecknoglobals
+	{},
+}
+
 func (c *Conn) fragmentHandshake(dtlsHandshake *handshake.Handshake) ([][]byte, error) {
+	messageSize := dtlsHandshake.Message.MarshalSize()
+	numFragments := (messageSize-1)/c.maximumTransmissionUnit + 1
+
+	fragmentedHandshakes := make([][]byte, numFragments)
+
+	if numFragments == 1 {
+		fragmentedHandshake := make([]byte, handshake.HeaderLength+messageSize)
+
+		headerFragment := handshake.Header{
+			Type:            dtlsHandshake.Header.Type,
+			Length:          dtlsHandshake.Header.Length,
+			MessageSequence: dtlsHandshake.Header.MessageSequence,
+			FragmentOffset:  uint32(0),
+			FragmentLength:  uint32(messageSize), //nolint:gosec // G115
+		}
+
+		_, err := headerFragment.MarshalTo(fragmentedHandshake)
+		if err != nil {
+			return nil, err
+		}
+
+		_, err = dtlsHandshake.Message.MarshalTo(fragmentedHandshake[handshake.HeaderLength:])
+		if err != nil {
+			return nil, err
+		}
+
+		fragmentedHandshakes[0] = fragmentedHandshake
+
+		return fragmentedHandshakes, nil
+	}
+
 	content, err := dtlsHandshake.Message.Marshal()
 	if err != nil {
 		return nil, err
 	}
 
-	fragmentedHandshakes := make([][]byte, 0)
-
 	contentFragments := util.SplitBytes(content, c.maximumTransmissionUnit)
 	if len(contentFragments) == 0 {
-		contentFragments = [][]byte{
-			{},
-		}
+		contentFragments = noContentFragments
 	}
 
 	offset := 0
-	for _, contentFragment := range contentFragments {
+	for i, contentFragment := range contentFragments {
 		contentFragmentLen := len(contentFragment)
 
-		headerFragment := &handshake.Header{
+		headerFragment := handshake.Header{
 			Type:            dtlsHandshake.Header.Type,
 			Length:          dtlsHandshake.Header.Length,
 			MessageSequence: dtlsHandshake.Header.MessageSequence,
 			FragmentOffset:  uint32(offset),
 			FragmentLength:  uint32(contentFragmentLen), //nolint:gosec // G115
 		}
 
+		fragmentedHandshake := make([]byte, handshake.HeaderLength+contentFragmentLen)
+
 		offset += contentFragmentLen
 
-		fragmentedHandshake, err := headerFragment.Marshal()
+		_, err := headerFragment.MarshalTo(fragmentedHandshake)
 		if err != nil {
 			return nil, err
 		}
 
-		fragmentedHandshake = append(fragmentedHandshake, contentFragment...)
-		fragmentedHandshakes = append(fragmentedHandshakes, fragmentedHandshake)
+		copy(fragmentedHandshake[handshake.HeaderLength:], contentFragment)
+
+		fragmentedHandshakes[i] = fragmentedHandshake
 	}
 
 	return fragmentedHandshakes, nil
@@ -2756,12 +2822,7 @@ func (c *Conn) close(byUser bool) error {
 }
 
 func (c *Conn) isConnectionClosed() bool {
-	select {
-	case <-c.closed.Done():
-		return true
-	default:
-		return false
-	}
+	return c.closed.Err() != nil
 }
 
 func (c *Conn) setLocalEpoch(epoch uint16) {

internal/closer/closer.go

@@ -6,6 +6,7 @@ package closer
 
 import (
 	"context"
+	"time"
 )
 
 // Closer allows for each signaling a channel for shutdown.
@@ -48,3 +49,11 @@ func (c *Closer) Err() error {
 func (c *Closer) Close() {
 	c.closeFunc()
 }
+
+func (c *Closer) Deadline() (deadline time.Time, ok bool) {
+	return c.ctx.Deadline()
+}
+
+func (c *Closer) Value(key any) any {
+	return c.ctx.Value(key)
+}

internal/flight/cache.go

@@ -5,7 +5,6 @@
 package flight
 
 import (
-	"bytes"
 	"fmt"
 	"sync"
 
@@ -45,7 +44,7 @@ func (h *Cache) Push(data []byte, epoch, messageSequence uint16, typ handshake.T
 	h.mu.Lock()
 	defer h.mu.Unlock()
 
-	h.cache = append(h.cache, &HandshakeCacheItem{Data: bytes.Clone(data), Epoch: epoch, MessageSequence: messageSequence, Typ: typ, IsClient: isClient})
+	h.cache = append(h.cache, &HandshakeCacheItem{Data: data, Epoch: epoch, MessageSequence: messageSequence, Typ: typ, IsClient: isClient})
 }
 
 // PullExact returns the handshake message with the requested sequence and

internal/handshake/fsm12.go

@@ -5,6 +5,7 @@ package dtlshandshake
 
 import (
 	"context"
+	"sync"
 	"time"
 
 	dtlsconfig "github.com/pion/dtls/v3/internal/config"
@@ -147,8 +148,23 @@ func (s *fsm12) send(ctx context.Context, c Conn) (State, error) {
 	return StateWaiting, nil
 }
 
+var timerPool = sync.Pool{ //nolint:gochecknoglobals
+	New: func() any {
+		timer := time.NewTimer(1 * time.Hour)
+		timer.Stop()
+
+		return timer
+	},
+}
+
 func (s *fsm12) wait(ctx context.Context, conn Conn) (State, error) { //nolint:gocognit,cyclop
-	retransmitTimer := time.NewTimer(s.retransmitInterval)
+	retransmitTimer := timerPool.Get().(*time.Timer) //nolint:forcetypeassert
+	retransmitTimer.Reset(s.retransmitInterval)
+	defer func() {
+		retransmitTimer.Stop()
+		timerPool.Put(retransmitTimer)
+	}()
+
 	for {
 		select {
 		case state := <-conn.RecvHandshake():

pkg/protocol/handshake/handshake.go

@@ -131,19 +131,18 @@ func (h *Handshake) Marshal() ([]byte, error) {
 		return nil, dtlserrors.ErrUnableToMarshalFragmented
 	}
 
-	message, err := h.Message.Marshal()
+	out := make([]byte, HeaderLength+h.Message.MarshalSize())
+	n, err := h.Message.MarshalTo(out[HeaderLength:])
 	if err != nil {
 		return nil, err
 	}
 
-	out := make([]byte, HeaderLength+len(message))
-	h.Header.Length = uint32(len(message)) //nolint:gosec // handshake messages are bounded to uint24 on the wire.
+	h.Header.Length = uint32(n) //nolint:gosec // handshake messages are bounded to uint24 on the wire.
 	h.Header.FragmentLength = h.Header.Length
 	h.Header.Type = h.Message.Type()
 	if _, err = h.Header.MarshalTo(out); err != nil {
 		return nil, err
 	}
-	copy(out[HeaderLength:], message)
 
 	return out, nil
 }