Add missing files, update style

internal/pcap/reorder.go

@@ -0,0 +1,216 @@
+package pcap
+
+import (
+	"cmp"
+	"errors"
+	"io"
+	"slices"
+	"sync"
+	"sync/atomic"
+)
+
+// A reorderEntry is 16 bytes, which bounds the memory needed per packet.
+type reorderEntry struct {
+	ts  uint64 // seconds << 32 | fraction (µs or ns, consistent within a file)
+	pos uint64 // record offset << lenBits | record length (including the 16 byte record header)
+}
+
+const (
+	lenBits   = 28 // fits MaxRecordLen + 16
+	maxOffset = 1<<(64-lenBits) - 1
+)
+
+func (e reorderEntry) off() int64 { return int64(e.pos >> lenBits) }
+func (e reorderEntry) len() int   { return int(e.pos & (1<<lenBits - 1)) }
+
+// Tuning for WriteTo: packets are copied in batches of up to batchBytes,
+// using parallel reads since random reads from SSDs are much faster with
+// several requests in flight. The next batch is read while the previous one
+// is written.
+const (
+	batchPackets = 16384
+	readers      = 16
+)
+
+var batchBytes = 8 << 20 // variable for tests
+
+// Reorder sorts the packets of a pcap file by timestamp. NewReorder reads
+// and indexes the file, so all parse errors are reported before WriteTo
+// produces any output.
+type Reorder struct {
+	r          io.ReaderAt
+	header     []byte
+	entries    []reorderEntry
+	end        int64 // offset just after the last complete record
+	outOfOrder int
+	truncated  bool
+}
+
+// NewReorder indexes the pcap file in r (of the given size) and sorts the
+// index by timestamp. The sort is stable: packets with equal timestamps keep
+// their original order. It needs 16 bytes of memory per packet.
+func NewReorder(r io.ReaderAt, size int64) (*Reorder, error) {
+	pr, err := NewReader(io.NewSectionReader(r, 0, size))
+	if err != nil {
+		return nil, err
+	}
+	ro := &Reorder{r: r, header: slices.Clone(pr.Header())}
+	off := int64(len(ro.header))
+	var latest uint64
+	for {
+		rec, err := pr.Next()
+		if err == io.EOF {
+			break
+		}
+		if errors.Is(err, ErrTruncated) {
+			ro.truncated = true
+			break
+		}
+		if err != nil {
+			return nil, err
+		}
+		if off > maxOffset {
+			return nil, errors.New("file too large to reorder")
+		}
+		ts := uint64(pr.order.Uint32(rec[0:4]))<<32 | uint64(pr.order.Uint32(rec[4:8]))
+		if ts < latest {
+			ro.outOfOrder++
+		}
+		latest = max(latest, ts)
+		ro.entries = append(ro.entries, reorderEntry{ts: ts, pos: uint64(off)<<lenBits | uint64(len(rec))})
+		off += int64(len(rec))
+
+		if len(ro.entries) == 1024 {
+			// Size the index from the average packet size so far, to avoid
+			// repeatedly growing (and copying) a large slice.
+			avg := (off - int64(len(ro.header))) / 1024
+			est := size / avg
+			ro.entries = slices.Grow(ro.entries, int(est+est/16))
+		}
+	}
+	if len(ro.entries) == 0 {
+		return nil, errors.New("the pcap file contains no packets")
+	}
+	ro.end = off
+	ro.entries = slices.Clip(ro.entries)
+	if ro.outOfOrder > 0 {
+		slices.SortStableFunc(ro.entries, func(a, b reorderEntry) int { return cmp.Compare(a.ts, b.ts) })
+	}
+	return ro, nil
+}
+
+// Packets returns the number of packets.
+func (ro *Reorder) Packets() int { return len(ro.entries) }
+
+// OutOfOrder returns the number of packets with a timestamp earlier than a
+// packet before them in the input.
+func (ro *Reorder) OutOfOrder() int { return ro.outOfOrder }
+
+// Truncated reports whether the input ended with an incomplete packet
+// record, which is left out of the output.
+func (ro *Reorder) Truncated() bool { return ro.truncated }
+
+type batch struct {
+	buf []byte
+	err error
+}
+
+// WriteTo writes the file header followed by the packets in timestamp order.
+func (ro *Reorder) WriteTo(w io.Writer) (int64, error) {
+	if ro.outOfOrder == 0 {
+		return io.Copy(w, io.NewSectionReader(ro.r, 0, ro.end))
+	}
+
+	n, err := w.Write(ro.header)
+	total := int64(n)
+	if err != nil {
+		return total, err
+	}
+
+	// A producer goroutine reads batches into two alternating buffers while
+	// this goroutine writes them out.
+	full := make(chan batch)
+	free := make(chan []byte, 2)
+	free <- nil
+	free <- nil
+	done := make(chan struct{})
+	var wg sync.WaitGroup
+	defer func() {
+		close(done)
+		wg.Wait()
+	}()
+	wg.Go(func() {
+		defer close(full)
+		for i := 0; i < len(ro.entries); {
+			var buf []byte
+			select {
+			case buf = <-free:
+			case <-done:
+				return
+			}
+			var err error
+			i, buf, err = ro.readBatch(i, buf)
+			select {
+			case full <- batch{buf, err}:
+			case <-done:
+				return
+			}
+			if err != nil {
+				return
+			}
+		}
+	})
+
+	for b := range full {
+		if b.err != nil {
+			return total, b.err
+		}
+		n, err := w.Write(b.buf)
+		total += int64(n)
+		if err != nil {
+			return total, err
+		}
+		free <- b.buf
+	}
+	return total, nil
+}
+
+// readBatch reads the records of entries[i:j] into buf, choosing j so that
+// the batch stays within batchBytes and batchPackets (but holds at least one
+// record). It returns j and the filled buffer.
+func (ro *Reorder) readBatch(i int, buf []byte) (int, []byte, error) {
+	j, size := i, 0
+	for j < len(ro.entries) && j-i < batchPackets && (j == i || size+ro.entries[j].len() <= batchBytes) {
+		size += ro.entries[j].len()
+		j++
+	}
+	entries := ro.entries[i:j]
+	dst := make([]int, len(entries)+1)
+	for k, e := range entries {
+		dst[k+1] = dst[k] + e.len()
+	}
+	buf = slices.Grow(buf[:0], size)[:size]
+
+	var (
+		next atomic.Int64
+		wg   sync.WaitGroup
+		errs [readers]error
+	)
+	for r := range min(readers, len(entries)) {
+		wg.Go(func() {
+			for {
+				k := int(next.Add(1) - 1)
+				if k >= len(entries) {
+					return
+				}
+				e := entries[k]
+				if _, err := ro.r.ReadAt(buf[dst[k]:dst[k+1]], e.off()); err != nil {
+					errs[r] = err
+					return
+				}
+			}
+		})
+	}
+	wg.Wait()
+	return j, buf, errors.Join(errs[:]...)
+}