mirror of
https://github.com/freemkv/libfreemkv.git
synced 2026-10-07 18:13:10 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e635c9556f | ||
|
|
d7b5c30f5d | ||
|
|
739a276a39 | ||
|
|
5f1028a62a | ||
|
|
ea15d212de | ||
|
|
00673c8ec3 | ||
|
|
bfefb4cb5b | ||
|
|
1b95193517 | ||
|
|
7dcac44136 | ||
|
|
4da559e39f | ||
|
|
eeca250f69 | ||
|
|
c51b3181f2 | ||
|
|
2a31a47434 | ||
|
|
b08662f95c | ||
|
|
4594196b6e | ||
|
|
9d6974df3e | ||
|
|
3394a5b3fe | ||
|
|
e2c7e20329 | ||
|
|
ff22385a61 | ||
|
|
d39d3933f5 | ||
|
|
31dc3d35c3 | ||
|
|
f7165a2020 | ||
|
|
d836f8b1e7 | ||
|
|
8919de5b03 | ||
|
|
0dc1f12108 | ||
|
|
6abc25a7e8 | ||
|
|
4d22cf632f | ||
|
|
611ad77586 | ||
|
|
60590193e3 | ||
|
|
5f24843d3c | ||
|
|
ecf9a4d01e | ||
|
|
8ec767339e | ||
|
|
335e8b68cc | ||
|
|
1137b4ea1f | ||
|
|
7788c560d1 | ||
|
|
84e0c2ca69 | ||
|
|
f4a881b25d | ||
|
|
1edda7fedb | ||
|
|
3dd71a0650 | ||
|
|
579c94bd23 | ||
|
|
3b2a68fe7a | ||
|
|
5b56f112a0 | ||
|
|
f27e4c3088 | ||
|
|
b3576f2ae6 | ||
|
|
39f2d0e992 | ||
|
|
50867516b8 | ||
|
|
193e7c1680 | ||
|
|
b6463729e8 | ||
|
|
dffee56102 | ||
|
|
d8f27cdee0 | ||
|
|
fc13268dc6 | ||
|
|
d9b55f02d0 | ||
|
|
35d66d2059 | ||
|
|
d703ce439b | ||
|
|
6112a26d15 | ||
|
|
2a33364253 | ||
|
|
fa3872ddcc | ||
|
|
c427389f36 | ||
|
|
6980f562df | ||
|
|
e3b2c9d850 | ||
|
|
a383200ef1 | ||
|
|
51d15f551e | ||
|
|
523af461a5 | ||
|
|
695b65149e | ||
|
|
d3f9560689 | ||
|
|
1fa5a7d27f | ||
|
|
55219070f1 | ||
|
|
a34c419521 | ||
|
|
b179846f5d | ||
|
|
e110e80e6e | ||
|
|
597fa34099 | ||
|
|
52eeb949e3 | ||
|
|
6ca62b8411 | ||
|
|
04195c27d8 | ||
|
|
5a8f8e54e1 | ||
|
|
1a2c830cf6 | ||
|
|
637c2bfe36 | ||
|
|
0d28357b4d | ||
|
|
5495332f07 | ||
|
|
e22fc6fd47 |
@@ -4,3 +4,4 @@ Cargo.lock
|
||||
*.swo
|
||||
.DS_Store
|
||||
.cargo/
|
||||
.claude/worktrees/
|
||||
|
||||
+227
@@ -1,5 +1,232 @@
|
||||
# Changelog
|
||||
|
||||
## 0.25.9 (2026-05-20)
|
||||
|
||||
### Fixed
|
||||
|
||||
- **MKB record-type identification.** `mkb_find_mk_dv` and `mkb_version`
|
||||
had their record-type constants swapped. The Verify Media Key record
|
||||
is Type 0x81 (AACS 1.0) or Type 0x86 (AACS 2.0/2.1), not 0x10. The
|
||||
Type-and-Version record is Type 0x10, not 0x81. PK and DK derivation
|
||||
paths therefore failed silently on every disc, masking how often the
|
||||
fallback paths could have succeeded. Fixed; mk_dv extracted at
|
||||
offset 4 of the verify record, MKB version read at offset 8 of the
|
||||
Type 0x10 record body. Tests added for both forms.
|
||||
|
||||
### Added
|
||||
|
||||
- **Built-in AACS keys.** Four device keys (covering MKB v01-v82+) and
|
||||
three processing keys (covering v63-v68) are now compiled into
|
||||
`libfreemkv` directly. DVDs and Blu-rays (AACS 1.0) decrypt with
|
||||
zero external files. Combined with the existing 31 CSS player keys
|
||||
in `css/auth.rs`, the library is self-sufficient for all DVD and
|
||||
AACS 1.0 content.
|
||||
- **Operator plugin slot.** `~/.config/freemkv/local_keys.cfg` is
|
||||
loaded automatically (same format as `keydb.cfg`) and layered
|
||||
additively on top of the built-ins and the main keydb.cfg. For
|
||||
operators who derive their own keys, this is the additive surface
|
||||
to drop them into — main keydb.cfg from upstream auto-update stays
|
||||
separate and overwritable.
|
||||
- `KeyDb::with_builtins()`, `KeyDb::load_or_builtins()`,
|
||||
`KeyDb::empty()` constructors.
|
||||
|
||||
### Changed
|
||||
|
||||
- `Disc::scan` no longer errors when keydb.cfg is absent. With
|
||||
built-ins covering DVD/BD, missing keydb is treated as "no UHD
|
||||
keys available" — the AACS 2.0 / UHD code path surfaces a
|
||||
specific error only when the disc actually needs keys that
|
||||
aren't in built-ins, main keydb, or local plugin.
|
||||
|
||||
## 0.25.8 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.8. No libfreemkv code changes —
|
||||
unified versioning bump only. autorip 0.25.7 image was DOA
|
||||
(missing +x on the binary); see autorip changelog.
|
||||
|
||||
## 0.25.7 (2026-05-20)
|
||||
|
||||
### Fixed — BU40N firmware wedge on KEYDB-miss discs
|
||||
|
||||
`disc::encrypt::do_handshake` used to fire up to 16 AACS authenticate
|
||||
attempts back-to-back with no pause between them. Each attempt is
|
||||
5-10 SCSI REPORT_KEY/SEND_KEY exchanges. On a disc whose host cert
|
||||
isn't in our KEYDB (or one the drive rejects), that's 80-160 SCSI
|
||||
commands hammered at the drive in a few hundred milliseconds — and
|
||||
the BU40N (plus most consumer optical drives) responds by entering
|
||||
a fast-fail firmware wedge state where every subsequent CDB returns
|
||||
`ILLEGAL_REQUEST/INVALID_FIELD_IN_CDB` (sense 0x05/0x24) until the
|
||||
drive is physically power-cycled.
|
||||
|
||||
Live wedge event on rip1 2026-05-20 during a Barbie UHD scan
|
||||
(KEYDB miss) confirmed the diagnosis and motivated this fix.
|
||||
|
||||
Defence-in-depth:
|
||||
- `MAX_CERT_ATTEMPTS` capped at 3 (was 16). If three different
|
||||
host certs all fail, more won't help — the drive doesn't have a
|
||||
match.
|
||||
- 1-second sleep between attempts. Gives the drive's firmware time
|
||||
to recover internal state between auth challenges.
|
||||
- Bail immediately on any sense code with sense_key = 0x05
|
||||
(`ILLEGAL_REQUEST`). The drive isn't merely rejecting our cert
|
||||
— it's saying "I won't talk to you anymore" — so trying more
|
||||
certs would deepen the wedge.
|
||||
|
||||
## 0.25.6 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.6 release. No behavioural changes.
|
||||
|
||||
## 0.25.5 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.5 release. No behavioural changes.
|
||||
|
||||
## 0.25.4 (2026-05-20)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.4 release. No behavioural changes
|
||||
in the library — the v0.25.4 release ships the autorip-side
|
||||
in-container NFS mount option.
|
||||
|
||||
## 0.25.3 (2026-05-19)
|
||||
|
||||
### Changed
|
||||
|
||||
- Synchronized to autorip 0.25.3 release. No behavioural changes
|
||||
in the library — the v0.25.3 release ships the autorip-side
|
||||
parallel mux pipeline.
|
||||
|
||||
## 0.25.2 (2026-05-19)
|
||||
|
||||
### Fixed
|
||||
|
||||
- **DTS-HD codec ID** — `MkvTrack::audio` now emits `A_DTS/MA` for
|
||||
DTS-HD MA and `A_DTS/HR` for DTS-HD HR instead of mislabelling
|
||||
both as plain `A_DTS`. Strict players (Plex transcoder, some
|
||||
hardware decoders, AV receivers) reject lossless DTS-HD MA
|
||||
payload when the track advertises `A_DTS` because the ID
|
||||
implies the 1.5 Mbps core-only bitstream.
|
||||
- **PGS subtitle BlockDuration** — the PGS parser is now stateful:
|
||||
it pairs each display PCS with the following empty/clear PCS to
|
||||
compute a duration, and the MKV muxer emits a `BlockGroup` +
|
||||
`BlockDuration` for subtitles that carry one. Without this the
|
||||
last bitmap lingered on screen until the next display set
|
||||
replaced it (or until end of file).
|
||||
|
||||
### Changed
|
||||
|
||||
- `codec::Frame` gains `duration_ns: Option<u64>` (set by parsers
|
||||
that can compute one; currently only PGS).
|
||||
- `pes::PesFrame` gains `duration_ns: Option<u64>` (in-memory only;
|
||||
not part of the on-wire serialization).
|
||||
- `MkvMuxer::write_frame` now takes a final `duration_ns: Option<u64>`
|
||||
parameter. When `Some`, the frame is emitted as a `BlockGroup`
|
||||
with `BlockDuration` instead of a `SimpleBlock`.
|
||||
|
||||
## 0.25.1 (2026-05-19)
|
||||
|
||||
### New — autorip event_fn plumbing
|
||||
|
||||
- `PrefetchedSectorSource::new_with_events(reader, extents, batch,
|
||||
halt, event_fn)` — same producer-thread pipeline as `new()`, plus
|
||||
an optional callback fired from the producer with `BytesRead`
|
||||
events after every successful batch. Lets the autorip multipass
|
||||
+ resume mux paths drive their progress UI from the highway
|
||||
without polling the consumer side.
|
||||
- `build_iso_pipeline` gains an `event_fn` parameter so callers can
|
||||
wire the same callback through one ctor.
|
||||
|
||||
### Changed
|
||||
|
||||
- `pes::Stream` trait: new default `errors() -> u64` method
|
||||
(default returns 0). Lets `Box<dyn Stream>` callers query the
|
||||
skip-on-error counter without downcasting. `DiscStream` overrides
|
||||
to surface its existing `errors` field.
|
||||
|
||||
### Removed (breaking)
|
||||
|
||||
- `DiscStream::new_pipeline` and `DiscStream::read_pipeline`
|
||||
deleted. All file-backed mux now uses `build_iso_pipeline` →
|
||||
`PipelinedPesStream`. `DiscStream` is now the single-threaded
|
||||
inline reader used by autorip's live-drive single-pass path only.
|
||||
- `DiscStream::demux_thread` and `demux_rx` fields removed.
|
||||
- `M2tsStream::open` deleted; `Mode::Read` variant removed.
|
||||
`m2ts://` URLs go through the internal `build_m2ts_pipeline`
|
||||
helper in `mux/resolve.rs` → `PipelinedPesStream`. `M2tsStream`
|
||||
is a write-only sink now.
|
||||
|
||||
## 0.25.0 (2026-05-19)
|
||||
|
||||
### New — the freemkv mux throughput "highway"
|
||||
|
||||
Three-stage pipelined PES read path. Read+decrypt runs on a
|
||||
producer thread, M2TS demux runs on a second thread, codec parse
|
||||
runs on the caller's thread. Communication between stages is via
|
||||
bounded `crossbeam_channel` with a recycled buffer pool — no
|
||||
allocations or memcpys in the steady-state hot loop.
|
||||
|
||||
**Throughput on the rip1 testbed (Civil_War UHD, 62 GiB ISO →
|
||||
`null://`, single-thread caller):**
|
||||
|
||||
| | MB/s |
|
||||
| ------------------------------------- | -----: |
|
||||
| 0.23.2 baseline | 60 |
|
||||
| + memchr SIMD HEVC start-code scan | 69 |
|
||||
| + ts.feed no-copy boundary | 72 |
|
||||
| + `PrefetchedSectorSource` (producer) | 124 |
|
||||
| + `DemuxThread` (3-stage pipeline) | 135 |
|
||||
| + zero-copy recycled buffer pool | 148 |
|
||||
| + 16 KiB initial PesAssembler buffer | 162 |
|
||||
| + mimalloc allocator in freemkv CLI | 200+ |
|
||||
| **+ warm cache** | **660** |
|
||||
|
||||
The new public API:
|
||||
|
||||
- `libfreemkv::PrefetchedSectorSource` — wraps any `SectorSource`,
|
||||
spawns a producer thread, exposes recycled-buffer channels.
|
||||
- `libfreemkv::io::byte_prefetcher::BytePrefetcher` —
|
||||
`std::io::Read` analogue for byte-stream sources (m2ts files,
|
||||
sockets, stdin).
|
||||
- `libfreemkv::mux::demux_thread::DemuxThread` — the M2TS demux
|
||||
worker; `spawn_zero_copy` takes either prefetcher's channels.
|
||||
- `libfreemkv::PipelinedPesStream` — the read-side `Stream` impl
|
||||
that runs codec parse on the caller thread.
|
||||
- `libfreemkv::build_iso_pipeline(reader, title, keys, batch,
|
||||
format, halt)` — the canonical ctor that wires all three stages
|
||||
for an ISO file source.
|
||||
|
||||
### Changed (breaking)
|
||||
|
||||
- `IsoSectorReader` (the naive duplicate of `FileSectorSource`)
|
||||
deleted. `FileSectorSource` is the sole file-backed sector
|
||||
source; it carries the SEQUENTIAL fadvise hint, the periodic
|
||||
DONTNEED page-cache eviction, and (new) the per-read
|
||||
`readahead()` async-prefetch syscall.
|
||||
- `mux::input("iso://...", &opts)` now returns a
|
||||
`PipelinedPesStream`. Function signature
|
||||
`input(&str, &InputOptions) -> io::Result<Box<dyn Stream>>` is
|
||||
unchanged; callers that treated the return value as
|
||||
`Box<dyn Stream>` keep working.
|
||||
|
||||
### Other
|
||||
|
||||
- `FileSectorSource` exposes per-OS `prefetch()` hooks
|
||||
(`readahead(2)` on Linux, `fcntl(F_RDADVISE)` on macOS, no-op on
|
||||
Windows + other).
|
||||
- AACS decrypt thread-pool: env var renamed from
|
||||
`FREEMKV_DECRYPT_THREADS` to `FREEMKV_THREADS`; default raised
|
||||
to `cores.clamp(1, 64)`.
|
||||
- `PesAssembler` initial buffer capacity 256 KiB → 16 KiB (avoids
|
||||
the 64-page first-touch fault tax on every PES boundary).
|
||||
- HEVC / H.264 `find_start_code` swapped onto
|
||||
`memchr::memmem::find` (SIMD).
|
||||
|
||||
## 0.18.4 (2026-05-09)
|
||||
|
||||
### Build / CI hardening — no library code changes
|
||||
|
||||
+18
-1
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "libfreemkv"
|
||||
version = "0.20.8"
|
||||
version = "0.25.9"
|
||||
edition = "2024"
|
||||
rust-version = "1.86"
|
||||
license = "AGPL-3.0-only"
|
||||
@@ -28,6 +28,23 @@ base64 = "0.22.1"
|
||||
# under project docs ("Acceptable strings: debug/trace logging"). Consumers (autorip)
|
||||
# wire a tracing subscriber and pipe events into the JSONL debug log.
|
||||
tracing = "0.1"
|
||||
# Bounded MPSC channel with kernel-wakeup send_timeout. Used by `io::pipeline`
|
||||
# so the halt-aware send/finish loops can BLOCK on consumer drain instead of
|
||||
# polling — the 50 ms poll cadence of the previous mpsc-based impl capped mux
|
||||
# throughput at ~1 MB/s (see (internal)/memory/
|
||||
# feedback_send_with_halt_poll_throttle.md, 0.21.7).
|
||||
crossbeam-channel = "0.5"
|
||||
# Persistent work-stealing thread pool for parallel AACS unit
|
||||
# decryption. Per-call std::thread::scope spawned fresh OS threads
|
||||
# and that overhead dominated for typical batch sizes (60 units).
|
||||
# rayon's global pool initialises once on first use.
|
||||
rayon = "1"
|
||||
# SIMD-accelerated bytestring search. Drives the HEVC/H.264 start-code
|
||||
# scan in `mux::codec::h264::find_start_code` — naive byte-by-byte
|
||||
# walk is ~500 MB/s single-thread on x86_64; memchr's vectorised
|
||||
# `memmem::find` for the 3-byte `00 00 01` needle hits ~5 GB/s on
|
||||
# AVX2-capable hosts.
|
||||
memchr = "2"
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
libc = "0.2"
|
||||
|
||||
@@ -6,6 +6,8 @@
|
||||
|
||||
Rust library for 4K UHD / Blu-ray / DVD optical drives. Drive access, disc scanning, stream labels, AACS decryption, CSS decryption, KEYDB updates, and content reading in one crate. Bundled drive profiles — no external files needed.
|
||||
|
||||
Built-in keys cover DVDs and Blu-rays (AACS 1.0). For UHD (AACS 2.0 / 2.1) discs, an optional `keydb.cfg` supplies disc-specific volume unique keys.
|
||||
|
||||
**12+ MB/s** sustained read speeds on BD. Full init: unlock, firmware upload, speed calibration — all from pure Rust.
|
||||
|
||||
Multi-lingual by design — the library outputs structured data and numeric error codes, never English text. Build any UI or localization on top.
|
||||
@@ -18,7 +20,7 @@ Part of the [freemkv](https://github.com/freemkv) project.
|
||||
|
||||
```toml
|
||||
[dependencies]
|
||||
libfreemkv = "0.18"
|
||||
libfreemkv = "0.25"
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
@@ -56,8 +58,7 @@ output.finish()?;
|
||||
For damaged discs the library exposes two flat verbs — `Disc::sweep` for the
|
||||
forward Pass 1 and `Disc::patch` for retrying bad ranges. The library never
|
||||
loops; the multipass policy is the caller's job. See
|
||||
[`docs/rip-recovery.md`](docs/rip-recovery.md) and the design notes in
|
||||
`(internal)/memory/0_18_redesign.md`.
|
||||
[`docs/rip-recovery.md`](docs/rip-recovery.md).
|
||||
|
||||
```rust
|
||||
use libfreemkv::{SweepOptions, PatchOptions};
|
||||
@@ -122,7 +123,13 @@ loop {
|
||||
|
||||
Streams implement `FrameSource` (read) and/or `FrameSink` (write); direction is type-checked. `input()` / `output()` resolve URL strings to PES stream instances. All URLs use the `scheme://path` format — bare paths are rejected.
|
||||
|
||||
AACS decryption requires a KEYDB.cfg file. If available at `~/.config/aacs/KEYDB.cfg` or passed via `ScanOptions`, the library handles everything — handshake, key derivation, and per-sector decryption — without the application needing to know anything about encryption.
|
||||
### Keys
|
||||
|
||||
DVDs (CSS) and Blu-rays (AACS 1.0) decrypt out of the box using bundled material — public AACS keys covering MKB version ranges are compiled into the library.
|
||||
|
||||
UHD discs (AACS 2.0 / 2.1) use per-disc volume unique keys, so they need a `keydb.cfg`. If one is available at `~/.config/freemkv/keydb.cfg` (or passed via `ScanOptions`), libfreemkv handles handshake, key resolution, and per-sector decryption transparently.
|
||||
|
||||
Operators with additional device keys, processing keys, or VUKs they have derived themselves can drop them into `~/.config/freemkv/local_keys.cfg`. This file uses the same syntax as `keydb.cfg` and is loaded on top of the built-in keys and any main `keydb.cfg` — entries are additive.
|
||||
|
||||
## Architecture
|
||||
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
//! Built-in public AACS 1.0 keys.
|
||||
//!
|
||||
//! Compile-time tables of the device keys (DK) and processing keys (PK)
|
||||
//! required for AACS 1.0 MKB processing. These values are well-known
|
||||
//! public AACS inputs that cover the MKB version ranges shipped on
|
||||
//! retail Blu-ray / UHD discs.
|
||||
//!
|
||||
//! With these built-ins, libfreemkv can resolve AACS 1.0 encryption for
|
||||
//! any disc whose VUK can be derived from MKB + device-key / processing-key
|
||||
//! paths — no external keydb.cfg file is required. Operators who want to
|
||||
//! supply additional keys (for example, future AACS 2.x derivations) can
|
||||
//! drop a `local_keys.cfg` into `$HOME/.config/freemkv/` in the same
|
||||
//! format as `keydb.cfg`; see [`crate::aacs::KeyDb::load_or_builtins`].
|
||||
|
||||
use super::keydb::DeviceKey;
|
||||
|
||||
/// A built-in device key entry. Mirrors [`DeviceKey`] but stored as a
|
||||
/// `const`-friendly POD type with an additional MKB range tag for
|
||||
/// diagnostics. Convert via [`BuiltinDeviceKey::to_device_key`].
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub(crate) struct BuiltinDeviceKey {
|
||||
pub key: [u8; 16],
|
||||
pub device_node: u16,
|
||||
pub key_uv: u32,
|
||||
pub u_mask_shift: u8,
|
||||
/// MKB version range tag (for logging / diagnostics only).
|
||||
#[allow(dead_code)]
|
||||
pub mkb_range: &'static str,
|
||||
}
|
||||
|
||||
impl BuiltinDeviceKey {
|
||||
pub(crate) fn to_device_key(self) -> DeviceKey {
|
||||
DeviceKey {
|
||||
key: self.key,
|
||||
node: self.device_node,
|
||||
uv: self.key_uv,
|
||||
u_mask_shift: self.u_mask_shift,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Public AACS 1.0 device keys covering MKB versions v01 through v82+.
|
||||
///
|
||||
/// Each entry contributes a subset-difference path through the MKB tree,
|
||||
/// so the four together cover the MKB ranges shipped on retail Blu-ray
|
||||
/// and UHD discs to date.
|
||||
pub(crate) const BUILTIN_DEVICE_KEYS: &[BuiltinDeviceKey] = &[
|
||||
BuiltinDeviceKey {
|
||||
key: [
|
||||
0x5F, 0xB8, 0x6E, 0xF1, 0x27, 0xC1, 0x9C, 0x17, 0x1E, 0x79, 0x9F, 0x61, 0xC2, 0x7B,
|
||||
0xDC, 0x2A,
|
||||
],
|
||||
device_node: 0x0800,
|
||||
key_uv: 0x0000_0400,
|
||||
u_mask_shift: 0x17,
|
||||
mkb_range: "v01-v48",
|
||||
},
|
||||
BuiltinDeviceKey {
|
||||
key: [
|
||||
0x38, 0x84, 0x16, 0x73, 0xE2, 0xB4, 0xE0, 0x51, 0x91, 0x65, 0x98, 0x99, 0x60, 0x6C,
|
||||
0xFF, 0xB8,
|
||||
],
|
||||
device_node: 0x0C00,
|
||||
key_uv: 0x0000_0A00,
|
||||
u_mask_shift: 0x0B,
|
||||
mkb_range: "v49-v71",
|
||||
},
|
||||
BuiltinDeviceKey {
|
||||
key: [
|
||||
0x86, 0x1B, 0x37, 0x19, 0xB0, 0x2F, 0x24, 0xBE, 0x6F, 0x1A, 0x30, 0xE2, 0xE3, 0xAB,
|
||||
0xEE, 0x94,
|
||||
],
|
||||
device_node: 0x0C40,
|
||||
key_uv: 0x0000_0D00,
|
||||
u_mask_shift: 0x0A,
|
||||
mkb_range: "v72+",
|
||||
},
|
||||
BuiltinDeviceKey {
|
||||
key: [
|
||||
0x7C, 0x06, 0xDE, 0xAE, 0x7F, 0x49, 0xB5, 0x51, 0xDA, 0xF5, 0x38, 0xC8, 0xCF, 0x18,
|
||||
0x11, 0xC9,
|
||||
],
|
||||
device_node: 0x0E20,
|
||||
key_uv: 0x0000_0E23,
|
||||
u_mask_shift: 0x02,
|
||||
mkb_range: "v82+",
|
||||
},
|
||||
];
|
||||
|
||||
/// Public AACS 1.0 processing keys for specific MKB versions.
|
||||
///
|
||||
/// Each value is a precomputed media-key-precursor that resolves MKB
|
||||
/// processing for the version range called out beside it. Provided as a
|
||||
/// fast path so an exhaustive device-key MKB walk is not required when a
|
||||
/// matching PK is available.
|
||||
pub(crate) const BUILTIN_PROCESSING_KEYS: &[[u8; 16]] = &[
|
||||
// v63
|
||||
[
|
||||
0x76, 0xDD, 0xD7, 0x09, 0x32, 0x16, 0xD2, 0x8C, 0x15, 0x04, 0x9A, 0x6B, 0x9C, 0x5C, 0x18,
|
||||
0xB9,
|
||||
],
|
||||
// v64-v65
|
||||
[
|
||||
0x3B, 0x32, 0x3C, 0x7A, 0x9A, 0xFC, 0x09, 0x21, 0x83, 0x1D, 0x24, 0x72, 0x39, 0x82, 0x3D,
|
||||
0xE6,
|
||||
],
|
||||
// v66-v68
|
||||
[
|
||||
0x7A, 0x4F, 0x40, 0xD8, 0x69, 0x6B, 0x7B, 0x15, 0x9B, 0xE8, 0x17, 0x6C, 0xC9, 0xED, 0xB8,
|
||||
0x5C,
|
||||
],
|
||||
];
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn builtin_device_keys_count() {
|
||||
assert_eq!(BUILTIN_DEVICE_KEYS.len(), 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builtin_processing_keys_count() {
|
||||
assert_eq!(BUILTIN_PROCESSING_KEYS.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_builtin_device_key_value() {
|
||||
let expected: [u8; 16] = [
|
||||
0x5F, 0xB8, 0x6E, 0xF1, 0x27, 0xC1, 0x9C, 0x17, 0x1E, 0x79, 0x9F, 0x61, 0xC2, 0x7B,
|
||||
0xDC, 0x2A,
|
||||
];
|
||||
assert_eq!(BUILTIN_DEVICE_KEYS[0].key, expected);
|
||||
assert_eq!(BUILTIN_DEVICE_KEYS[0].device_node, 0x0800);
|
||||
assert_eq!(BUILTIN_DEVICE_KEYS[0].key_uv, 0x0000_0400);
|
||||
assert_eq!(BUILTIN_DEVICE_KEYS[0].u_mask_shift, 0x17);
|
||||
}
|
||||
}
|
||||
+272
-3
@@ -54,6 +54,16 @@ pub struct DiscEntry {
|
||||
pub unit_keys: Vec<(u32, [u8; 16])>,
|
||||
}
|
||||
|
||||
/// Path to the operator-managed local plugin file. Operators drop
|
||||
/// additional AACS keys here (same on-disk format as keydb.cfg) and
|
||||
/// they layer transparently on top of the built-ins and the main
|
||||
/// keydb.cfg at load time. Returns `None` if `HOME` (or `USERPROFILE`
|
||||
/// on Windows) is not set in the environment.
|
||||
pub fn local_plugin_path() -> Option<std::path::PathBuf> {
|
||||
let home = std::env::var_os("HOME").or_else(|| std::env::var_os("USERPROFILE"))?;
|
||||
Some(std::path::PathBuf::from(home).join(".config/freemkv/local_keys.cfg"))
|
||||
}
|
||||
|
||||
/// Parse a hex string like "0xABCD..." into bytes.
|
||||
pub(crate) fn parse_hex(s: &str) -> Option<Vec<u8>> {
|
||||
let s = s.trim().trim_start_matches("0x").trim_start_matches("0X");
|
||||
@@ -89,6 +99,74 @@ pub(crate) fn parse_hex20(s: &str) -> Option<[u8; 20]> {
|
||||
}
|
||||
|
||||
impl KeyDb {
|
||||
/// Construct an empty KeyDb.
|
||||
pub fn empty() -> Self {
|
||||
KeyDb {
|
||||
device_keys: Vec::new(),
|
||||
processing_keys: Vec::new(),
|
||||
host_certs: Vec::new(),
|
||||
disc_entries: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Construct a KeyDb pre-populated with the compiled-in public AACS 1.0
|
||||
/// device keys and processing keys. Sufficient on its own to derive
|
||||
/// VUKs for any AACS 1.0 disc whose MKB version is covered.
|
||||
pub fn with_builtins() -> Self {
|
||||
let mut db = Self::empty();
|
||||
db.add_builtins();
|
||||
db
|
||||
}
|
||||
|
||||
/// Push the built-in AACS 1.0 device keys and processing keys into
|
||||
/// this KeyDb. Duplicates (identical device-key triples or identical
|
||||
/// processing-key bytes) are silently skipped — first-seen wins.
|
||||
fn add_builtins(&mut self) {
|
||||
for entry in crate::aacs::builtin_keys::BUILTIN_DEVICE_KEYS {
|
||||
self.add_device_key_dedup(entry.to_device_key());
|
||||
}
|
||||
for pk in crate::aacs::builtin_keys::BUILTIN_PROCESSING_KEYS {
|
||||
self.add_processing_key_dedup(*pk);
|
||||
}
|
||||
}
|
||||
|
||||
fn add_device_key_dedup(&mut self, dk: DeviceKey) {
|
||||
let dup = self
|
||||
.device_keys
|
||||
.iter()
|
||||
.any(|x| x.node == dk.node && x.uv == dk.uv && x.u_mask_shift == dk.u_mask_shift);
|
||||
if !dup {
|
||||
self.device_keys.push(dk);
|
||||
}
|
||||
}
|
||||
|
||||
fn add_processing_key_dedup(&mut self, pk: [u8; 16]) {
|
||||
if !self.processing_keys.iter().any(|x| *x == pk) {
|
||||
self.processing_keys.push(pk);
|
||||
}
|
||||
}
|
||||
|
||||
/// Merge another KeyDb's entries into this one, additively. Existing
|
||||
/// entries are kept; duplicates from `other` are silently dropped.
|
||||
/// Used to layer external keydb.cfg / local plugin contents on top of
|
||||
/// the built-ins.
|
||||
pub fn merge_from(&mut self, other: KeyDb) {
|
||||
for dk in other.device_keys {
|
||||
self.add_device_key_dedup(dk);
|
||||
}
|
||||
for pk in other.processing_keys {
|
||||
self.add_processing_key_dedup(pk);
|
||||
}
|
||||
// Host certs and disc entries do not have a stable "identity"
|
||||
// key suitable for dedup beyond byte-equality, so append host
|
||||
// certs as-is and insert disc entries with first-wins semantics
|
||||
// by hash.
|
||||
self.host_certs.extend(other.host_certs);
|
||||
for (hash, entry) in other.disc_entries {
|
||||
self.disc_entries.entry(hash).or_insert(entry);
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a KEYDB.cfg file from a string.
|
||||
pub fn parse(data: &str) -> Self {
|
||||
let mut db = KeyDb {
|
||||
@@ -152,10 +230,47 @@ impl KeyDb {
|
||||
db
|
||||
}
|
||||
|
||||
/// Load KEYDB.cfg from a file path.
|
||||
/// Load a KEYDB.cfg from disk, layered on top of the compiled-in
|
||||
/// built-in keys. If `path` does not exist, returns a KeyDb that
|
||||
/// contains only the built-ins (no error). If `path` exists but
|
||||
/// cannot be read, the underlying I/O error is returned.
|
||||
///
|
||||
/// An operator plugin slot at `$HOME/.config/freemkv/local_keys.cfg`
|
||||
/// (same on-disk format) is also layered on top when present. This
|
||||
/// lets operators drop additional keys at runtime without editing
|
||||
/// the main keydb.cfg.
|
||||
pub fn load(path: &std::path::Path) -> std::io::Result<Self> {
|
||||
let data = std::fs::read_to_string(path)?;
|
||||
Ok(Self::parse(&data))
|
||||
let mut db = Self::with_builtins();
|
||||
if path.exists() {
|
||||
let data = std::fs::read_to_string(path)?;
|
||||
db.merge_from(Self::parse(&data));
|
||||
}
|
||||
db.merge_local_plugin();
|
||||
Ok(db)
|
||||
}
|
||||
|
||||
/// Load only the built-in keys plus the operator plugin slot at
|
||||
/// `$HOME/.config/freemkv/local_keys.cfg` (if present). Used when
|
||||
/// no main keydb.cfg path is configured.
|
||||
pub fn load_or_builtins() -> Self {
|
||||
let mut db = Self::with_builtins();
|
||||
db.merge_local_plugin();
|
||||
db
|
||||
}
|
||||
|
||||
/// If `$HOME/.config/freemkv/local_keys.cfg` exists, layer its
|
||||
/// contents on top of this KeyDb. Errors reading the plugin file
|
||||
/// are silently ignored — the plugin is a best-effort augmentation.
|
||||
fn merge_local_plugin(&mut self) {
|
||||
let Some(path) = local_plugin_path() else {
|
||||
return;
|
||||
};
|
||||
if !path.exists() {
|
||||
return;
|
||||
}
|
||||
if let Ok(data) = std::fs::read_to_string(&path) {
|
||||
self.merge_from(Self::parse(&data));
|
||||
}
|
||||
}
|
||||
|
||||
/// Look up a disc by its hash. Returns the VUK if found.
|
||||
@@ -391,6 +506,160 @@ mod tests {
|
||||
assert_eq!(hc.certificate.len(), 92);
|
||||
}
|
||||
|
||||
// Mutex serializes tests that mutate process-wide environment
|
||||
// variables (HOME). Tests in the same module run on threads by
|
||||
// default; a shared lock here prevents env-var bleed-through.
|
||||
use std::sync::Mutex;
|
||||
static ENV_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
#[test]
|
||||
fn keydb_default_has_builtins() {
|
||||
let db = KeyDb::with_builtins();
|
||||
assert_eq!(db.device_keys.len(), 4);
|
||||
assert_eq!(db.processing_keys.len(), 3);
|
||||
assert!(db.host_certs.is_empty());
|
||||
assert!(db.disc_entries.is_empty());
|
||||
|
||||
let expected_first: [u8; 16] = [
|
||||
0x5F, 0xB8, 0x6E, 0xF1, 0x27, 0xC1, 0x9C, 0x17, 0x1E, 0x79, 0x9F, 0x61, 0xC2, 0x7B,
|
||||
0xDC, 0x2A,
|
||||
];
|
||||
assert_eq!(db.device_keys[0].key, expected_first);
|
||||
assert_eq!(db.device_keys[0].node, 0x0800);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keydb_load_layers_on_builtins() {
|
||||
// Synthetic keydb.cfg with one extra device key not in built-ins.
|
||||
let extra_dk = "| DK | DEVICE_KEY 0xAABBCCDDEEFF00112233445566778899 | DEVICE_NODE 0x1234 | KEY_UV 0x00001234 | KEY_U_MASK_SHIFT 0x05 ; extra\n";
|
||||
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let tmp = tempdir_isolated_home();
|
||||
let cfg_path = tmp.path().join("keydb.cfg");
|
||||
std::fs::write(&cfg_path, extra_dk).unwrap();
|
||||
|
||||
let db = KeyDb::load(&cfg_path).unwrap();
|
||||
// 4 built-ins + 1 extra = 5 DKs total
|
||||
assert_eq!(db.device_keys.len(), 5);
|
||||
// Built-ins are still present and come first
|
||||
assert_eq!(db.device_keys[0].node, 0x0800);
|
||||
// Extra entry is layered on top
|
||||
assert!(db.device_keys.iter().any(|d| d.node == 0x1234));
|
||||
assert_eq!(db.processing_keys.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keydb_load_missing_file_falls_back_to_builtins() {
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let _tmp = tempdir_isolated_home();
|
||||
let db = KeyDb::load(std::path::Path::new("/nonexistent/keydb.cfg")).unwrap();
|
||||
assert_eq!(db.device_keys.len(), 4);
|
||||
assert_eq!(db.processing_keys.len(), 3);
|
||||
assert!(db.host_certs.is_empty());
|
||||
assert!(db.disc_entries.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keydb_local_plugin_layered() {
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let tmp = tempdir_isolated_home();
|
||||
|
||||
// Drop a local_keys.cfg into the synthetic HOME with one extra DK.
|
||||
let plugin_dir = tmp.path().join(".config/freemkv");
|
||||
std::fs::create_dir_all(&plugin_dir).unwrap();
|
||||
let plugin_path = plugin_dir.join("local_keys.cfg");
|
||||
std::fs::write(
|
||||
&plugin_path,
|
||||
"| DK | DEVICE_KEY 0x112233445566778899AABBCCDDEEFF00 | DEVICE_NODE 0xABCD | KEY_UV 0x0000ABCD | KEY_U_MASK_SHIFT 0x07 ; plugin\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
// load_or_builtins must include the plugin entry.
|
||||
let db = KeyDb::load_or_builtins();
|
||||
assert_eq!(db.device_keys.len(), 5);
|
||||
assert!(db.device_keys.iter().any(|d| d.node == 0xABCD));
|
||||
|
||||
// load() of a non-existent main keydb must also include the plugin.
|
||||
let db2 = KeyDb::load(std::path::Path::new("/nonexistent/keydb.cfg")).unwrap();
|
||||
assert!(db2.device_keys.iter().any(|d| d.node == 0xABCD));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_with_no_keydb_path() {
|
||||
// Higher-level resolution: when ScanOptions has no keydb_path
|
||||
// and no keydb.cfg is in standard search paths, the loader
|
||||
// returns a KeyDb populated with the built-ins (plus any
|
||||
// plugin entries). This is the "AACS 1.0 just works" path.
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let _tmp = tempdir_isolated_home();
|
||||
// Wipe XDG/system paths from view by pointing HOME at an empty
|
||||
// dir and confirming there's nothing in our plugin slot. Then
|
||||
// assert the no-path entry point returns built-ins only.
|
||||
let db = KeyDb::load_or_builtins();
|
||||
assert_eq!(db.device_keys.len(), 4);
|
||||
assert_eq!(db.processing_keys.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keydb_dedup_keeps_first() {
|
||||
// Loading the same DK twice (built-in + identical entry in cfg)
|
||||
// results in one entry, not two.
|
||||
let dup_dk = "| DK | DEVICE_KEY ***REMOVED*** | DEVICE_NODE 0x0800 | KEY_UV 0x00000400 | KEY_U_MASK_SHIFT 0x17 ; duplicate of builtin\n";
|
||||
let _guard = ENV_LOCK.lock().unwrap();
|
||||
let tmp = tempdir_isolated_home();
|
||||
let cfg_path = tmp.path().join("keydb.cfg");
|
||||
std::fs::write(&cfg_path, dup_dk).unwrap();
|
||||
|
||||
let db = KeyDb::load(&cfg_path).unwrap();
|
||||
// Still 4, not 5 — the duplicate was deduped
|
||||
assert_eq!(db.device_keys.len(), 4);
|
||||
}
|
||||
|
||||
/// Build a temporary directory and point `HOME`/`USERPROFILE` at it
|
||||
/// so the local-plugin loader sees an isolated, empty environment by
|
||||
/// default. The returned [`TempDir`] auto-cleans on drop.
|
||||
fn tempdir_isolated_home() -> TempDir {
|
||||
let tmp = TempDir::new();
|
||||
// SAFETY: tests serialize via ENV_LOCK before calling this.
|
||||
unsafe {
|
||||
std::env::set_var("HOME", tmp.path());
|
||||
std::env::set_var("USERPROFILE", tmp.path());
|
||||
}
|
||||
tmp
|
||||
}
|
||||
|
||||
/// Minimal scoped temp directory — auto-removes on drop.
|
||||
/// Avoids adding the `tempfile` crate as a dependency.
|
||||
struct TempDir {
|
||||
path: std::path::PathBuf,
|
||||
}
|
||||
impl TempDir {
|
||||
fn new() -> Self {
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
let nanos = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_nanos())
|
||||
.unwrap_or(0);
|
||||
let tid = std::thread::current().id();
|
||||
let path = std::env::temp_dir().join(format!(
|
||||
"libfreemkv-keydb-test-{:?}-{}-{}",
|
||||
tid,
|
||||
std::process::id(),
|
||||
nanos
|
||||
));
|
||||
std::fs::create_dir_all(&path).expect("create tempdir");
|
||||
TempDir { path }
|
||||
}
|
||||
fn path(&self) -> &std::path::Path {
|
||||
&self.path
|
||||
}
|
||||
}
|
||||
impl Drop for TempDir {
|
||||
fn drop(&mut self) {
|
||||
let _ = std::fs::remove_dir_all(&self.path);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_full_keydb() {
|
||||
let path = match keydb_path() {
|
||||
|
||||
+70
-30
@@ -169,8 +169,9 @@ pub fn parse_unit_key_ro(data: &[u8], aacs2: bool) -> Option<UnitKeyFile> {
|
||||
/// This is the fast path — no subset-difference tree traversal needed.
|
||||
///
|
||||
/// MKB format:
|
||||
/// Record type 0x10 = Verify Media Key Record (has mk_dv)
|
||||
/// Record type 0x81 = Type and Version Record (has MKB version)
|
||||
/// Record type 0x10 = Type and Version Record (has MKB version)
|
||||
/// Record type 0x81 = Verify Media Key Record, AACS 1.0 (has mk_dv)
|
||||
/// Record type 0x86 = Verify Media Key Record, AACS 2.0/2.1 (has mk_dv)
|
||||
/// Record type 0x04 = Subset-Difference Index (has UVS entries)
|
||||
/// Record type 0x07 = Explicit Subset-Difference Record (has cvalues)
|
||||
pub fn derive_media_key_from_pk(mkb: &[u8], processing_keys: &[[u8; 16]]) -> Option<[u8; 16]> {
|
||||
@@ -244,10 +245,10 @@ fn validate_processing_key(
|
||||
None
|
||||
}
|
||||
|
||||
/// Find Verify Media Key Record (type 0x10) in MKB.
|
||||
/// Find Verify Media Key Record (type 0x81 for AACS 1.0, 0x86 for AACS 2.0/2.1) in MKB.
|
||||
fn mkb_find_mk_dv(mkb: &[u8]) -> Option<[u8; 16]> {
|
||||
let mut pos = 0;
|
||||
let mut type10_seen: Vec<(usize, usize)> = Vec::new();
|
||||
let mut verify_rec_seen: Vec<(u8, usize, usize)> = Vec::new();
|
||||
while pos + 4 <= mkb.len() {
|
||||
let rec_type = mkb[pos];
|
||||
let rec_len = u32::from_be_bytes([0, mkb[pos + 1], mkb[pos + 2], mkb[pos + 3]]) as usize;
|
||||
@@ -255,17 +256,18 @@ fn mkb_find_mk_dv(mkb: &[u8]) -> Option<[u8; 16]> {
|
||||
break;
|
||||
}
|
||||
|
||||
if rec_type == 0x10 {
|
||||
type10_seen.push((pos, rec_len));
|
||||
if rec_type == 0x81 || rec_type == 0x86 {
|
||||
verify_rec_seen.push((rec_type, pos, rec_len));
|
||||
}
|
||||
|
||||
if rec_type == 0x10 && rec_len >= 20 {
|
||||
// mk_dv is at offset 4 (after record header)
|
||||
if (rec_type == 0x81 || rec_type == 0x86) && rec_len >= 20 {
|
||||
// mk_dv is at offset 4 of the record (after the 4-byte header)
|
||||
let mut dv = [0u8; 16];
|
||||
dv.copy_from_slice(&mkb[pos + 4..pos + 20]);
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "mkb_mk_dv_found",
|
||||
rec_type,
|
||||
pos,
|
||||
rec_len,
|
||||
"mk_dv extracted from MKB"
|
||||
@@ -277,9 +279,9 @@ fn mkb_find_mk_dv(mkb: &[u8]) -> Option<[u8; 16]> {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "mkb_mk_dv_not_found",
|
||||
type10_seen = ?type10_seen,
|
||||
verify_rec_seen = ?verify_rec_seen,
|
||||
scanned_bytes = pos,
|
||||
"no 0x10 record with rec_len>=20 found"
|
||||
"no 0x81/0x86 record with rec_len>=20 found"
|
||||
);
|
||||
None
|
||||
}
|
||||
@@ -320,7 +322,8 @@ fn mkb_find_cvalues(mkb: &[u8]) -> Option<Vec<u8>> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Get MKB version from Type and Version Record (type 0x81).
|
||||
/// Get MKB version from Type and Version Record (type 0x10).
|
||||
/// Version is a BE u32 at offset 8 of the record body (offset 12 from `pos`).
|
||||
pub fn mkb_version(mkb: &[u8]) -> Option<u32> {
|
||||
let mut pos = 0;
|
||||
while pos + 4 <= mkb.len() {
|
||||
@@ -330,12 +333,12 @@ pub fn mkb_version(mkb: &[u8]) -> Option<u32> {
|
||||
break;
|
||||
}
|
||||
|
||||
if rec_type == 0x81 && rec_len >= 8 {
|
||||
if rec_type == 0x10 && rec_len >= 12 {
|
||||
return Some(u32::from_be_bytes([
|
||||
mkb[pos + 4],
|
||||
mkb[pos + 5],
|
||||
mkb[pos + 6],
|
||||
mkb[pos + 7],
|
||||
mkb[pos + 8],
|
||||
mkb[pos + 9],
|
||||
mkb[pos + 10],
|
||||
mkb[pos + 11],
|
||||
]));
|
||||
}
|
||||
pos += rec_len;
|
||||
@@ -935,23 +938,60 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mkb_version_parse() {
|
||||
// Synthetic MKB with Type and Version record (0x81)
|
||||
let mut mkb = vec![0u8; 32];
|
||||
// Record: type=0x81, length=12 (BE24)
|
||||
mkb[0] = 0x81;
|
||||
mkb[1] = 0x00;
|
||||
mkb[2] = 0x00;
|
||||
mkb[3] = 0x0C;
|
||||
// Version = 77
|
||||
mkb[4] = 0x00;
|
||||
mkb[5] = 0x00;
|
||||
mkb[6] = 0x00;
|
||||
mkb[7] = 77;
|
||||
|
||||
fn mkb_version_recognizes_type_0x10() {
|
||||
// Type-and-Version record: type=0x10, rec_len=12 (BE24).
|
||||
// Body is 8 bytes; the version u32 sits at offset 8 of the record.
|
||||
let mkb = [
|
||||
0x10, 0x00, 0x00, 0x0C, 0x48, 0x14, 0x10, 0x03, 0x00, 0x00, 0x00, 0x4D,
|
||||
];
|
||||
assert_eq!(mkb_version(&mkb), Some(77));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mkb_version_returns_none_on_empty() {
|
||||
assert_eq!(mkb_version(&[]), None);
|
||||
assert_eq!(mkb_version(&[0x10, 0x00]), None);
|
||||
// Type 0x10 record but rec_len < 12 → no version available.
|
||||
let short = [0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x01];
|
||||
assert_eq!(mkb_version(&short), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mkb_find_mk_dv_recognizes_type_0x81() {
|
||||
// First: type-0x10 type/version record (12 bytes), then type-0x81 verify record.
|
||||
// Verify record carries a known 16-byte mk_dv at offset 4 of the record body.
|
||||
let expected: [u8; 16] = [
|
||||
0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC, 0xDD, 0xEE,
|
||||
0xFF, 0x00,
|
||||
];
|
||||
let mut mkb = vec![
|
||||
0x10, 0x00, 0x00, 0x0C, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
|
||||
];
|
||||
// type=0x81, rec_len=24 (4-byte header + 16-byte mk_dv + 4-byte trailing zeros)
|
||||
mkb.extend_from_slice(&[0x81, 0x00, 0x00, 0x18]);
|
||||
mkb.extend_from_slice(&expected);
|
||||
mkb.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
|
||||
|
||||
assert_eq!(mkb_find_mk_dv(&mkb), Some(expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mkb_find_mk_dv_recognizes_type_0x86() {
|
||||
// AACS 2.0 form uses type 0x86 for the verify record.
|
||||
let expected: [u8; 16] = [
|
||||
0xDE, 0xAD, 0xBE, 0xEF, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A,
|
||||
0x0B, 0x0C,
|
||||
];
|
||||
let mut mkb = vec![
|
||||
0x10, 0x00, 0x00, 0x0C, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x4D,
|
||||
];
|
||||
mkb.extend_from_slice(&[0x86, 0x00, 0x00, 0x18]);
|
||||
mkb.extend_from_slice(&expected);
|
||||
mkb.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
|
||||
|
||||
assert_eq!(mkb_find_mk_dv(&mkb), Some(expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_keys_vuk_path() {
|
||||
// Test the full resolve chain using VUK path
|
||||
|
||||
@@ -13,6 +13,7 @@
|
||||
//! The VUK decrypts title keys from AACS/Unit_Key_RO.inf on disc.
|
||||
//! Title keys decrypt m2ts stream content (AES-128-CBC).
|
||||
|
||||
pub(crate) mod builtin_keys;
|
||||
pub mod decrypt;
|
||||
pub mod handshake;
|
||||
pub mod keydb;
|
||||
|
||||
+4
-3
@@ -4,13 +4,14 @@
|
||||
//! 1. Bus authentication (challenge-response) → bus key
|
||||
//! 2. Read disc key block (READ DVD STRUCTURE) → XOR with bus key → decrypt with player keys → disc key
|
||||
//! 3. Read title key (REPORT KEY format 0x04) → XOR with bus key → decrypt with disc key → title key
|
||||
//!
|
||||
//! Based on libdvdcss (VideoLAN) and Stevenson 1999 analysis.
|
||||
|
||||
use crate::drive::Drive;
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
// ── Player keys (from libdvdcss, Stevenson's PlayerKey cracker) ───────────
|
||||
// ── Built-in public DVD CSS player keys ────────────────────────────────────
|
||||
//
|
||||
// These 31 5-byte player keys are long-public CSS inputs. With them
|
||||
// compiled in, DVD ripping works with no external key file required.
|
||||
|
||||
const PLAYER_KEYS: [[u8; 5]; 31] = [
|
||||
[0x01, 0xaf, 0xe3, 0x12, 0x80],
|
||||
|
||||
+153
-16
@@ -3,9 +3,131 @@
|
||||
//! Decrypts sectors in-place using resolved keys from disc scanning.
|
||||
//! Handles AACS 1.0, AACS 2.0, and CSS transparently.
|
||||
//! The caller never sees encrypted data unless explicitly bypassed.
|
||||
//!
|
||||
//! ## Parallel AACS decrypt
|
||||
//!
|
||||
//! Each AACS aligned unit (6144 bytes) is decrypted INDEPENDENTLY of
|
||||
//! every other unit — per-unit key derivation from the unit_key plus
|
||||
//! the unit's own first-16-byte header. There is no cross-unit
|
||||
//! dependency, so a buffer of N units can be decrypted on N threads
|
||||
//! in parallel via a persistent rayon thread pool.
|
||||
//!
|
||||
//! Small buffers (< [`PARALLEL_MIN_UNITS`] units) fall through to the
|
||||
//! serial path to avoid pool dispatch overhead beating the per-unit
|
||||
//! AES work.
|
||||
//!
|
||||
//! ## Thread-count configuration — three layers
|
||||
//!
|
||||
//! Resolution order (highest wins):
|
||||
//! 1. The most recent [`set_decrypt_threads`] call with `n > 0`.
|
||||
//! Calling this *replaces* the live thread pool — useful for a
|
||||
//! settings-page slider in a long-running daemon.
|
||||
//! 2. `FREEMKV_THREADS` env var, if set and `> 0`. Single knob
|
||||
//! covering decrypt today, intended to also drive any future
|
||||
//! input-side / output-side worker pools.
|
||||
//! 3. Default: all available cores. Algorithm optimisation comes
|
||||
//! first — we measure single-thread performance to find serial
|
||||
//! bottlenecks before throwing parallelism at it — but once a
|
||||
//! pool is engaged we use the whole box. Hard cap at
|
||||
//! [`MAX_THREADS`] (rayon stack memory).
|
||||
|
||||
use crate::aacs;
|
||||
use crate::css;
|
||||
use rayon::prelude::*;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
/// Minimum units in a buffer before we pay the pool-dispatch cost of
|
||||
/// fanning out. Below this, serial is faster.
|
||||
const PARALLEL_MIN_UNITS: usize = 8;
|
||||
|
||||
/// Hard upper bound on configurable thread count. Anything larger is
|
||||
/// almost certainly a misconfiguration; rayon would happily allocate
|
||||
/// thousands of worker stacks otherwise.
|
||||
pub const MAX_THREADS: usize = 64;
|
||||
|
||||
/// Process-wide decrypt thread count override. `0` means "use env
|
||||
/// var, else default" — see [`decrypt_threads`] for the resolution
|
||||
/// order.
|
||||
static DECRYPT_THREADS: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
/// Current rayon pool. `RwLock<Option<Arc<...>>>` so that
|
||||
/// [`set_decrypt_threads`] can swap the pool out without leaking the
|
||||
/// old one and without blocking ongoing decrypt work (in-flight calls
|
||||
/// hold an `Arc` clone via [`decrypt_pool`] and finish on the old
|
||||
/// pool; new calls pick up the new pool).
|
||||
static DECRYPT_POOL: RwLock<Option<Arc<rayon::ThreadPool>>> = RwLock::new(None);
|
||||
|
||||
/// Configure how many threads to use for AACS unit decryption. A value
|
||||
/// of `0` resets to the env / default resolution. `1` forces serial.
|
||||
/// `N > 1` builds a new rayon pool of size N and atomically replaces
|
||||
/// the live pool.
|
||||
///
|
||||
/// Thread-safe. Live decrypt calls keep their previously-acquired
|
||||
/// pool reference for the rest of the call — no mid-call pool
|
||||
/// switch. Subsequent calls see the new pool.
|
||||
///
|
||||
/// Pool construction is ~ms-scale; safe to call from a settings POST
|
||||
/// handler.
|
||||
pub fn set_decrypt_threads(n: usize) {
|
||||
let clamped = n.min(MAX_THREADS);
|
||||
DECRYPT_THREADS.store(clamped, Ordering::Relaxed);
|
||||
// Drop the existing pool. Next decrypt_pool() call rebuilds with
|
||||
// the new resolved thread count.
|
||||
if let Ok(mut guard) = DECRYPT_POOL.write() {
|
||||
*guard = None;
|
||||
}
|
||||
}
|
||||
|
||||
/// Get (or lazily build) the active rayon thread pool. Returns an
|
||||
/// `Arc` so in-flight work survives a concurrent
|
||||
/// [`set_decrypt_threads`] swap.
|
||||
fn decrypt_pool() -> Arc<rayon::ThreadPool> {
|
||||
// Fast path: pool already built.
|
||||
if let Ok(guard) = DECRYPT_POOL.read() {
|
||||
if let Some(pool) = guard.as_ref() {
|
||||
return Arc::clone(pool);
|
||||
}
|
||||
}
|
||||
// Slow path: build a new one under the write lock. Double-check
|
||||
// after acquiring in case another caller built it first.
|
||||
let mut guard = DECRYPT_POOL.write().expect("DECRYPT_POOL RwLock poisoned");
|
||||
if let Some(pool) = guard.as_ref() {
|
||||
return Arc::clone(pool);
|
||||
}
|
||||
let n = decrypt_threads();
|
||||
let pool = Arc::new(
|
||||
rayon::ThreadPoolBuilder::new()
|
||||
.num_threads(n)
|
||||
.thread_name(|i| format!("freemkv-decrypt-{i}"))
|
||||
.build()
|
||||
.expect("rayon decrypt pool build failed"),
|
||||
);
|
||||
*guard = Some(Arc::clone(&pool));
|
||||
pool
|
||||
}
|
||||
|
||||
/// Current effective decrypt thread count. Resolution order:
|
||||
/// 1. Most recent [`set_decrypt_threads`] value (if > 0)
|
||||
/// 2. `FREEMKV_THREADS` env var (if set and > 0)
|
||||
/// 3. Default: all available cores, capped at [`MAX_THREADS`].
|
||||
pub fn decrypt_threads() -> usize {
|
||||
let explicit = DECRYPT_THREADS.load(Ordering::Relaxed);
|
||||
if explicit > 0 {
|
||||
return explicit;
|
||||
}
|
||||
let env = std::env::var("FREEMKV_THREADS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<usize>().ok())
|
||||
.unwrap_or(0);
|
||||
if env > 0 {
|
||||
return env.min(MAX_THREADS);
|
||||
}
|
||||
let cores = std::thread::available_parallelism()
|
||||
.map(|n| n.get())
|
||||
.unwrap_or(2);
|
||||
cores.clamp(1, MAX_THREADS)
|
||||
}
|
||||
|
||||
/// Resolved decryption state from disc scanning.
|
||||
/// Passed to `decrypt_sectors()` — the caller doesn't need to know
|
||||
@@ -57,29 +179,44 @@ pub fn decrypt_sectors(
|
||||
return Err(crate::error::Error::DecryptFailed);
|
||||
}
|
||||
};
|
||||
let rdk = read_data_key.as_ref();
|
||||
let rdk: Option<[u8; 16]> = *read_data_key;
|
||||
let unit_len = aacs::ALIGNED_UNIT_LEN;
|
||||
let nthreads = decrypt_threads();
|
||||
let chunks: Vec<&mut [u8]> = buf.chunks_mut(unit_len).collect();
|
||||
let nunits = chunks.len();
|
||||
|
||||
for chunk in buf.chunks_mut(unit_len) {
|
||||
// Per-unit decrypt closure. The is_unit_encrypted check is
|
||||
// a byte-0 heuristic; on a misfire we snapshot+restore via
|
||||
// the original bytes so non-m2ts (e.g. MPLS/CLPI nav files)
|
||||
// survive. See test `nav_file_unit_survives_decrypt_attempt`.
|
||||
let decrypt_one = |chunk: &mut [u8]| {
|
||||
if chunk.len() == unit_len && aacs::is_unit_encrypted(chunk) {
|
||||
// `is_unit_encrypted` is a byte-0 heuristic: it fires on any
|
||||
// unit whose first byte has the top 2 bits set, which is
|
||||
// correct for m2ts source packets (where those bits are the
|
||||
// copy-control marker) but false-positives on any other binary
|
||||
// data with similarly-shaped first bytes — notably MPLS/CLPI
|
||||
// navigation files that begin with ASCII magic ('M', 'H'…)
|
||||
// and survive sweep mixed in with encrypted m2ts payloads.
|
||||
// `decrypt_unit_full` self-checks via TS-sync verification and
|
||||
// returns false on a misfire, but it has already mutated the
|
||||
// chunk by then. Snapshot and restore on verification failure
|
||||
// — same pattern `decrypt_unit_try_keys` uses for multi-key
|
||||
// discs. Real m2ts units verify and stay decrypted; nav-file
|
||||
// sectors get scrambled briefly and then put back as-was.
|
||||
let original: Vec<u8> = chunk.to_vec();
|
||||
if !aacs::decrypt_unit_full(chunk, &uk, rdk) {
|
||||
if !aacs::decrypt_unit_full(chunk, &uk, rdk.as_ref()) {
|
||||
chunk.copy_from_slice(&original);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
if nthreads <= 1 || nunits < PARALLEL_MIN_UNITS {
|
||||
// Serial path: avoids thread-pool overhead for tiny
|
||||
// buffers; also the only path when caller pinned
|
||||
// single-threaded via FREEMKV_THREADS=1.
|
||||
for chunk in chunks {
|
||||
decrypt_one(chunk);
|
||||
}
|
||||
} else {
|
||||
// Parallel path via rayon's persistent global pool.
|
||||
// The pool is built once on first use (lazy_static-style)
|
||||
// and reused across every decrypt_sectors call — no
|
||||
// per-call OS thread spawn, no thread-creation latency
|
||||
// amortised per batch. Each unit decrypts independently
|
||||
// (own key derivation), so par_iter is sound.
|
||||
decrypt_pool().install(|| {
|
||||
chunks.into_par_iter().for_each(|chunk| {
|
||||
decrypt_one(chunk);
|
||||
});
|
||||
});
|
||||
}
|
||||
}
|
||||
DecryptKeys::Css { title_key } => {
|
||||
|
||||
+51
-7
@@ -61,9 +61,29 @@ impl Disc {
|
||||
"handshake starting"
|
||||
);
|
||||
|
||||
const MAX_CERT_ATTEMPTS: usize = 16;
|
||||
// v0.25.7 wedge fix. Pre-0.25.7 this loop fired up to 16 AACS
|
||||
// authenticate attempts back-to-back with no pause. Each attempt
|
||||
// is 5-10 SCSI REPORT_KEY/SEND_KEY exchanges. On a disc whose
|
||||
// host cert isn't in our KEYDB (or one the drive rejects),
|
||||
// that's 80-160 SCSI commands hammered at the drive in a
|
||||
// few hundred milliseconds — and the BU40N (and most consumer
|
||||
// optical drives) responds by entering a fast-fail firmware
|
||||
// wedge state where every subsequent CDB returns
|
||||
// ILLEGAL_REQUEST/INVALID_FIELD_IN_CDB (sense 05/24) until
|
||||
// power-cycled. Hit live on rip1 2026-05-20 during a Barbie
|
||||
// UHD scan: KEYDB miss → 16 cert attempts in a tight loop →
|
||||
// wedge → forced host reboot + drive disconnect to recover.
|
||||
//
|
||||
// Defense-in-depth: cap attempts, sleep between, and bail
|
||||
// early on the drive's wedge sense so any later regression
|
||||
// can't undo the protection silently.
|
||||
const MAX_CERT_ATTEMPTS: usize = 3;
|
||||
const PER_CERT_BACKOFF_MS: u64 = 1000;
|
||||
let mut last_err_code: Option<u16> = None;
|
||||
for (idx, hc) in keydb.host_certs.iter().take(MAX_CERT_ATTEMPTS).enumerate() {
|
||||
if idx > 0 {
|
||||
std::thread::sleep(std::time::Duration::from_millis(PER_CERT_BACKOFF_MS));
|
||||
}
|
||||
match aacs::handshake::aacs_authenticate(session, &hc.private_key, &hc.certificate) {
|
||||
Ok(mut auth) => {
|
||||
let volume_id = match aacs::handshake::read_volume_id(session, &mut auth) {
|
||||
@@ -94,7 +114,24 @@ impl Disc {
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
last_err_code = Some(e.code());
|
||||
let code = e.code();
|
||||
last_err_code = Some(code);
|
||||
// Drive wedge senses (any with high byte 0x05 =
|
||||
// ILLEGAL_REQUEST). The drive isn't merely
|
||||
// rejecting our cert — it's saying "I won't talk
|
||||
// to you anymore." Trying more certs makes the
|
||||
// wedge worse. Bail out immediately.
|
||||
let sense_key = ((code >> 8) & 0xFF) as u8;
|
||||
if sense_key == 0x05 {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "handshake_wedge_detected",
|
||||
cert_index = idx,
|
||||
error_code = code,
|
||||
"drive returned ILLEGAL_REQUEST during auth; bailing out to avoid wedge"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
}
|
||||
@@ -105,7 +142,8 @@ impl Disc {
|
||||
host_cert_count,
|
||||
tried = host_cert_count.min(MAX_CERT_ATTEMPTS),
|
||||
last_error_code = last_err_code,
|
||||
"all host certs in KEYDB rejected by drive"
|
||||
"all host certs in KEYDB rejected by drive (capped at {} attempts to prevent firmware wedge)",
|
||||
MAX_CERT_ATTEMPTS
|
||||
);
|
||||
// All host certs failed — return None, not a fake success
|
||||
None
|
||||
@@ -119,14 +157,20 @@ impl Disc {
|
||||
pub(super) fn resolve_encryption(
|
||||
udf_fs: &udf::UdfFs,
|
||||
reader: &mut dyn SectorSource,
|
||||
keydb_path: &std::path::Path,
|
||||
keydb_path: Option<&std::path::Path>,
|
||||
handshake: Option<&HandshakeResult>,
|
||||
) -> Result<AacsState> {
|
||||
use crate::aacs::{self, KeyDb};
|
||||
|
||||
let keydb = KeyDb::load(keydb_path).map_err(|_| Error::KeydbLoad {
|
||||
path: keydb_path.display().to_string(),
|
||||
})?;
|
||||
// Built-in AACS 1.0 keys are always available. When a keydb.cfg
|
||||
// path is supplied, layer it on top; otherwise fall back to
|
||||
// built-ins (plus the operator local-plugin slot, if any).
|
||||
let keydb = match keydb_path {
|
||||
Some(path) => KeyDb::load(path).map_err(|_| Error::KeydbLoad {
|
||||
path: path.display().to_string(),
|
||||
})?,
|
||||
None => KeyDb::load_or_builtins(),
|
||||
};
|
||||
|
||||
// Read AACS files from disc/image via UDF
|
||||
let uk_ro_data = udf_fs
|
||||
|
||||
+26
-31
@@ -1116,40 +1116,35 @@ impl Disc {
|
||||
udf_fs.find_dir("/AACS").is_some() || udf_fs.find_dir("/BDMV/AACS").is_some();
|
||||
|
||||
let (aacs, aacs_error) = if encrypted {
|
||||
match opts.resolve_keydb() {
|
||||
Some(keydb_path) => {
|
||||
match Self::resolve_encryption(&udf_fs, reader, &keydb_path, handshake.as_ref())
|
||||
{
|
||||
Ok(state) => (Some(state), None),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "scan_aacs_resolve_failed",
|
||||
error_code = e.code(),
|
||||
keydb = %keydb_path.display(),
|
||||
handshake_ok = handshake.is_some(),
|
||||
"AACS key resolution failed"
|
||||
);
|
||||
(None, Some(e))
|
||||
}
|
||||
}
|
||||
}
|
||||
None => {
|
||||
// KEYDB is now optional: the library ships built-in AACS 1.0
|
||||
// device + processing keys, so a missing keydb.cfg just falls
|
||||
// back to the built-ins (plus the operator local-plugin slot,
|
||||
// if any). External keydb.cfg layers on top when supplied.
|
||||
let keydb_path = opts.resolve_keydb();
|
||||
if keydb_path.is_none() {
|
||||
tracing::debug!(
|
||||
target: "freemkv::disc",
|
||||
phase = "scan_aacs_builtins_only",
|
||||
"no external KEYDB found; resolving with built-in AACS 1.0 keys"
|
||||
);
|
||||
}
|
||||
match Self::resolve_encryption(
|
||||
&udf_fs,
|
||||
reader,
|
||||
keydb_path.as_deref(),
|
||||
handshake.as_ref(),
|
||||
) {
|
||||
Ok(state) => (Some(state), None),
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
target: "freemkv::disc",
|
||||
phase = "scan_aacs_no_keydb",
|
||||
"encrypted disc but no KEYDB found in search paths"
|
||||
phase = "scan_aacs_resolve_failed",
|
||||
error_code = e.code(),
|
||||
keydb = ?keydb_path.as_ref().map(|p| p.display().to_string()),
|
||||
handshake_ok = handshake.is_some(),
|
||||
"AACS key resolution failed"
|
||||
);
|
||||
// Reuse KeydbLoad with sentinel path — adding a new Error
|
||||
// variant would be a breaking change for downstream
|
||||
// exhaustive matches. The path string makes the cause
|
||||
// unambiguous to autorip's message switch.
|
||||
(
|
||||
None,
|
||||
Some(crate::error::Error::KeydbLoad {
|
||||
path: String::from("<no keydb in search paths>"),
|
||||
}),
|
||||
)
|
||||
(None, Some(e))
|
||||
}
|
||||
}
|
||||
} else {
|
||||
|
||||
+14
@@ -71,6 +71,20 @@ impl Default for Halt {
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared poll interval for halt-aware loops.
|
||||
///
|
||||
/// `bounded_syscall` checks the cancellation flag and the deadline
|
||||
/// every [`POLL_INTERVAL`] while blocked on a worker; the same
|
||||
/// cadence governs `Pipeline::send_with_halt`'s `try_send` retry.
|
||||
/// 250 ms is the sweet spot between responsiveness (operator presses
|
||||
/// Stop, sees it take effect within ~quarter-second) and waste
|
||||
/// (atomic load + clock read is cheap but not free at thousands of
|
||||
/// hertz).
|
||||
///
|
||||
/// Centralised here so the half-dozen halt-polling loops across `io`
|
||||
/// can't drift apart silently.
|
||||
pub const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_millis(250);
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
+1
-9
@@ -54,15 +54,7 @@ use std::sync::mpsc::{RecvTimeoutError, sync_channel};
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::halt::Halt;
|
||||
|
||||
/// Granularity of the halt poll. The receive loop wakes every
|
||||
/// [`POLL_INTERVAL`] to (a) check the [`Halt`] token, then (b) check
|
||||
/// the overall deadline, then go back to waiting. 250 ms is a
|
||||
/// pragmatic balance: short enough that human-driven `/api/stop` feels
|
||||
/// responsive (< 0.5 s p99), long enough that the polling overhead is
|
||||
/// negligible against multi-second syscalls.
|
||||
const POLL_INTERVAL: Duration = Duration::from_millis(250);
|
||||
use crate::halt::{Halt, POLL_INTERVAL};
|
||||
|
||||
/// Failure outcome from a bounded syscall wrapper.
|
||||
#[derive(Debug)]
|
||||
|
||||
@@ -0,0 +1,330 @@
|
||||
//! Byte-sized bounded producer/consumer channel.
|
||||
//!
|
||||
//! Wraps `std::sync::mpsc::sync_channel` with a byte-accounting
|
||||
//! `Mutex<usize> + Condvar` cap. Sender blocks (cooperatively) when
|
||||
//! `used_bytes + item.byte_size() > capacity_bytes`. Receiver
|
||||
//! decrements `used_bytes` when it takes the item.
|
||||
//!
|
||||
//! Why: the existing producer→consumer channel between `DiscStream`
|
||||
//! (PES producer) and `MuxSink` (PES consumer) is bounded by frame
|
||||
//! count. Frame sizes vary 100× between metadata and keyframes, so a
|
||||
//! count-based cap either starves on small frames or buffers far too
|
||||
//! much memory on big ones. Byte-sized accounting sizes the buffer for
|
||||
//! the worst-case input stall (NFS read p99 ≈ 1–2 s × ~15 MB/s peak
|
||||
//! compressed bitrate ≈ ~30 MB) directly.
|
||||
//!
|
||||
//! The underlying mpsc channel is created with a very large slot count
|
||||
//! so the byte cap (not the slot count) is the real backpressure. Slot
|
||||
//! count is only there to give the kernel a small chunk to wake on.
|
||||
//!
|
||||
//! See `(internal)/memory/project_buffering_architecture.md` §
|
||||
//! Pipeline channel — sizing.
|
||||
|
||||
use std::sync::mpsc::{Receiver as MpscReceiver, RecvError, SendError, SyncSender, sync_channel};
|
||||
use std::sync::{Arc, Condvar, Mutex};
|
||||
|
||||
/// Default byte cap for the muxer's input channel. Sized to hide a
|
||||
/// worst-case ~2 s NFS read refill at UHD peak compressed bitrate
|
||||
/// (~15 MB/s); 64 MiB gives headroom. Tweakable; not magic.
|
||||
pub const BYTE_CHANNEL_DEFAULT_CAPACITY: usize = 64 * 1024 * 1024;
|
||||
|
||||
/// Slot capacity of the inner `sync_channel`. Large so the byte cap is
|
||||
/// the real backpressure mechanism — the mpsc slot count only exists
|
||||
/// to give the kernel a chunk to wake on. PES frames are typically
|
||||
/// ~700 B each, so 64 MiB ≈ 90 k frames; 200 k is comfortable headroom.
|
||||
const INNER_SLOT_CAPACITY: usize = 200_000;
|
||||
|
||||
/// Anything whose in-memory cost can be accounted by a single
|
||||
/// `usize`. Implement on the item type sent through [`Sender`].
|
||||
pub trait HasByteSize {
|
||||
/// Bytes this item contributes to the channel's used budget.
|
||||
/// Must be > 0 to make progress (a 0-byte item would never
|
||||
/// block the sender no matter the cap; see send_blocks_at_capacity
|
||||
/// test).
|
||||
fn byte_size(&self) -> usize;
|
||||
}
|
||||
|
||||
impl HasByteSize for crate::pes::PesFrame {
|
||||
fn byte_size(&self) -> usize {
|
||||
// Frame data + the fixed header overhead the serializer
|
||||
// writes (track + pts + keyframe + len). The `Vec<u8>` heap
|
||||
// allocation also has alloc-header overhead but that's
|
||||
// <0.1 % at typical frame sizes — folding it in would just
|
||||
// add noise to the budget.
|
||||
self.data.len() + 14
|
||||
}
|
||||
}
|
||||
|
||||
/// Shared book-keeping between [`Sender`] and [`Receiver`]. Wrapped in
|
||||
/// an `Arc` because both halves hold it independently.
|
||||
struct Accounting {
|
||||
used: Mutex<usize>,
|
||||
cv: Condvar,
|
||||
capacity: usize,
|
||||
}
|
||||
|
||||
/// Send half of the byte-bounded channel.
|
||||
///
|
||||
/// `send` blocks (on a `Condvar`) when adding the item would push
|
||||
/// `used_bytes` past `capacity_bytes`. Unblocks when the receiver
|
||||
/// `recv`s items out and notifies. Returns `Err(item)` if the
|
||||
/// receiver has been dropped — mirrors `mpsc::SyncSender::send`.
|
||||
pub struct Sender<T: HasByteSize> {
|
||||
tx: SyncSender<T>,
|
||||
acct: Arc<Accounting>,
|
||||
}
|
||||
|
||||
impl<T: HasByteSize> Clone for Sender<T> {
|
||||
fn clone(&self) -> Self {
|
||||
Sender {
|
||||
tx: self.tx.clone(),
|
||||
acct: self.acct.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: HasByteSize> Sender<T> {
|
||||
/// Push one item. Blocks until adding it would not exceed the
|
||||
/// capacity, then sends through the inner mpsc channel.
|
||||
pub fn send(&self, item: T) -> Result<(), SendError<T>> {
|
||||
let sz = item.byte_size();
|
||||
// Reserve capacity first. The reservation is observable to
|
||||
// other senders via `used`; only after we win the slot do we
|
||||
// hand the item to the inner mpsc channel. That ordering means
|
||||
// `used` is always a conservative upper bound on what's in the
|
||||
// mpsc queue + about-to-be-sent.
|
||||
{
|
||||
let mut used = self.acct.used.lock().expect("byte_channel poisoned");
|
||||
// An item bigger than the whole capacity will never fit; let
|
||||
// it through anyway as a one-shot reservation, otherwise the
|
||||
// sender deadlocks forever waiting for `used == 0` AND
|
||||
// nothing in flight. The receiver will drain it on the
|
||||
// other side. Same behaviour as `std::sync::mpsc` for
|
||||
// arbitrarily large messages.
|
||||
while *used + sz > self.acct.capacity && *used > 0 {
|
||||
used = self.acct.cv.wait(used).expect("byte_channel cv poisoned");
|
||||
}
|
||||
*used += sz;
|
||||
}
|
||||
match self.tx.send(item) {
|
||||
Ok(()) => Ok(()),
|
||||
Err(SendError(returned)) => {
|
||||
// Receiver dropped — refund the reservation so a later
|
||||
// sender on a clone doesn't observe phantom used bytes
|
||||
// (the receiver is gone so nobody will decrement).
|
||||
let mut used = self.acct.used.lock().expect("byte_channel poisoned");
|
||||
*used = used.saturating_sub(sz);
|
||||
self.acct.cv.notify_all();
|
||||
Err(SendError(returned))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Receive half of the byte-bounded channel.
|
||||
///
|
||||
/// `recv` blocks on the inner mpsc until an item is available, then
|
||||
/// decrements the byte-accounting and wakes any sender waiting on
|
||||
/// capacity.
|
||||
pub struct Receiver<T: HasByteSize> {
|
||||
rx: MpscReceiver<T>,
|
||||
acct: Arc<Accounting>,
|
||||
}
|
||||
|
||||
impl<T: HasByteSize> Receiver<T> {
|
||||
/// Take the next item. Returns `Err(RecvError)` when all senders
|
||||
/// have been dropped and the channel is empty.
|
||||
pub fn recv(&self) -> Result<T, RecvError> {
|
||||
let item = self.rx.recv()?;
|
||||
let sz = item.byte_size();
|
||||
let mut used = self.acct.used.lock().expect("byte_channel poisoned");
|
||||
*used = used.saturating_sub(sz);
|
||||
// Notify all so multi-sender setups wake every blocked sender,
|
||||
// not just one. Wasted wakeups are cheap; missed wakeups would
|
||||
// be a deadlock.
|
||||
self.acct.cv.notify_all();
|
||||
Ok(item)
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a byte-bounded channel with the given capacity in bytes.
|
||||
/// Returns a `(Sender, Receiver)` pair; clone the `Sender` for
|
||||
/// multi-producer setups.
|
||||
pub fn channel<T: HasByteSize>(capacity_bytes: usize) -> (Sender<T>, Receiver<T>) {
|
||||
let (tx, rx) = sync_channel::<T>(INNER_SLOT_CAPACITY);
|
||||
let acct = Arc::new(Accounting {
|
||||
used: Mutex::new(0),
|
||||
cv: Condvar::new(),
|
||||
capacity: capacity_bytes,
|
||||
});
|
||||
(
|
||||
Sender {
|
||||
tx,
|
||||
acct: acct.clone(),
|
||||
},
|
||||
Receiver { rx, acct },
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::thread;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Test payload — its `byte_size` returns whatever we passed at
|
||||
/// construction so capacity math is exact and predictable.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
struct Item {
|
||||
sz: usize,
|
||||
tag: u32,
|
||||
}
|
||||
|
||||
impl HasByteSize for Item {
|
||||
fn byte_size(&self) -> usize {
|
||||
self.sz
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn send_recv_round_trip() {
|
||||
let (tx, rx) = channel::<Item>(1024);
|
||||
for i in 0..5 {
|
||||
tx.send(Item { sz: 100, tag: i }).unwrap();
|
||||
}
|
||||
for i in 0..5 {
|
||||
let got = rx.recv().unwrap();
|
||||
assert_eq!(got, Item { sz: 100, tag: i });
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn byte_accounting_decrements_on_recv() {
|
||||
// Internal book-keeping check via observable side-effect: after
|
||||
// sending K items totalling N bytes and receiving them all, a
|
||||
// subsequent send of an N-byte item must NOT block (no items
|
||||
// in flight, all capacity refunded).
|
||||
let (tx, rx) = channel::<Item>(1024);
|
||||
for _ in 0..4 {
|
||||
tx.send(Item { sz: 256, tag: 0 }).unwrap();
|
||||
}
|
||||
for _ in 0..4 {
|
||||
rx.recv().unwrap();
|
||||
}
|
||||
// Cap is now fully available again. Send a 1024-byte item; the
|
||||
// `used > 0` guard means it goes through alone (no wait).
|
||||
let start = Instant::now();
|
||||
tx.send(Item { sz: 1024, tag: 99 }).unwrap();
|
||||
assert!(start.elapsed() < Duration::from_millis(100));
|
||||
let got = rx.recv().unwrap();
|
||||
assert_eq!(got.tag, 99);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn send_blocks_at_capacity_unblocks_on_recv() {
|
||||
// Cap = 200 bytes, item = 100 bytes. First two sends fit
|
||||
// exactly; the third must block until a recv frees capacity.
|
||||
let (tx, rx) = channel::<Item>(200);
|
||||
tx.send(Item { sz: 100, tag: 0 }).unwrap();
|
||||
tx.send(Item { sz: 100, tag: 1 }).unwrap();
|
||||
|
||||
let tx2 = tx.clone();
|
||||
let sent_at = Arc::new(Mutex::new(None::<Instant>));
|
||||
let sent_at2 = sent_at.clone();
|
||||
let h = thread::spawn(move || {
|
||||
tx2.send(Item { sz: 100, tag: 2 }).unwrap();
|
||||
*sent_at2.lock().unwrap() = Some(Instant::now());
|
||||
});
|
||||
|
||||
// Give the sender thread a head start; it should be parked in
|
||||
// `cv.wait` because used (200) + 100 > capacity (200).
|
||||
thread::sleep(Duration::from_millis(100));
|
||||
assert!(
|
||||
sent_at.lock().unwrap().is_none(),
|
||||
"third send should be blocked at capacity"
|
||||
);
|
||||
|
||||
// Drain one. Sender wakes and completes.
|
||||
let recv_at = Instant::now();
|
||||
let got = rx.recv().unwrap();
|
||||
assert_eq!(got.tag, 0);
|
||||
h.join().unwrap();
|
||||
|
||||
let sent_when = sent_at.lock().unwrap().unwrap();
|
||||
assert!(
|
||||
sent_when >= recv_at,
|
||||
"sender must complete AFTER receiver freed capacity"
|
||||
);
|
||||
|
||||
// Drain the remaining two.
|
||||
assert_eq!(rx.recv().unwrap().tag, 1);
|
||||
assert_eq!(rx.recv().unwrap().tag, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn item_larger_than_capacity_still_goes_through() {
|
||||
// Pathological case: a single item bigger than the capacity.
|
||||
// The guard `*used > 0` lets it through when the channel is
|
||||
// empty (otherwise the sender deadlocks forever). Matches
|
||||
// `mpsc::SyncSender` semantics for oversize messages.
|
||||
let (tx, rx) = channel::<Item>(100);
|
||||
tx.send(Item { sz: 1000, tag: 7 }).unwrap();
|
||||
let got = rx.recv().unwrap();
|
||||
assert_eq!(got, Item { sz: 1000, tag: 7 });
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn concurrent_send_recv_stress() {
|
||||
// 4 sender threads × 1k items each, 1 receiver. Verify byte
|
||||
// accounting stays sane (channel never deadlocks, every item
|
||||
// arrives exactly once) under contention.
|
||||
const SENDERS: u32 = 4;
|
||||
const PER_SENDER: u32 = 1000;
|
||||
const TOTAL: u32 = SENDERS * PER_SENDER;
|
||||
|
||||
let (tx, rx) = channel::<Item>(8 * 1024);
|
||||
let sent = Arc::new(AtomicUsize::new(0));
|
||||
let mut handles = Vec::new();
|
||||
for s in 0..SENDERS {
|
||||
let tx = tx.clone();
|
||||
let sent = sent.clone();
|
||||
handles.push(thread::spawn(move || {
|
||||
for i in 0..PER_SENDER {
|
||||
// Vary item size so accounting actually has to
|
||||
// multiplex differently-sized blockers. 1B → 256B.
|
||||
let sz = 1 + ((i as usize) % 256);
|
||||
tx.send(Item {
|
||||
sz,
|
||||
tag: s * PER_SENDER + i,
|
||||
})
|
||||
.unwrap();
|
||||
sent.fetch_add(1, Ordering::SeqCst);
|
||||
}
|
||||
}));
|
||||
}
|
||||
// Drop our local sender so the receiver can eventually see
|
||||
// RecvError once all sender clones are done. Cloning the
|
||||
// sender into each producer means each clone Drop'd separately.
|
||||
drop(tx);
|
||||
|
||||
let mut received = 0u32;
|
||||
while let Ok(_item) = rx.recv() {
|
||||
received += 1;
|
||||
}
|
||||
for h in handles {
|
||||
h.join().unwrap();
|
||||
}
|
||||
assert_eq!(received, TOTAL);
|
||||
assert_eq!(sent.load(Ordering::SeqCst) as u32, TOTAL);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn send_after_recv_dropped_returns_err() {
|
||||
let (tx, rx) = channel::<Item>(1024);
|
||||
drop(rx);
|
||||
let r = tx.send(Item { sz: 10, tag: 0 });
|
||||
assert!(r.is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,146 @@
|
||||
//! `BytePrefetcher` — `std::io::Read` analogue of
|
||||
//! [`crate::sector::PrefetchedSectorSource`].
|
||||
//!
|
||||
//! Spawns a producer thread that fills a bounded pool of `Vec<u8>`
|
||||
//! chunks from the underlying reader and ships them through a
|
||||
//! channel; the consumer pulls filled chunks, uses them, and sends
|
||||
//! the empty `Vec<u8>` back through a recycle channel so the
|
||||
//! producer can re-fill in place. Result: zero allocations and zero
|
||||
//! cross-thread frees in the steady-state hot loop.
|
||||
//!
|
||||
//! This is the byte-stream half of the freemkv mux highway —
|
||||
//! `BytePrefetcher` feeds [`crate::mux::demux_thread::DemuxThread`]
|
||||
//! for `m2ts://`, `network://`, `stdio://`, and any other stream
|
||||
//! whose source is an `io::Read` rather than a `SectorSource`.
|
||||
|
||||
use crate::halt::Halt;
|
||||
use crossbeam_channel::{Receiver, Sender, bounded};
|
||||
use std::io::Read;
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
/// Items flowing through the forward channel.
|
||||
pub type Batch = std::io::Result<Vec<u8>>;
|
||||
|
||||
/// Forward channel depth — how many filled buffers the producer can
|
||||
/// stay ahead by. Two is enough to absorb a moderate consumer stall
|
||||
/// without piling up bytes.
|
||||
const FORWARD_DEPTH: usize = 2;
|
||||
|
||||
/// Recycle channel depth = forward + 1 so the producer always has at
|
||||
/// least one buffer to fill while the consumer holds one.
|
||||
const RECYCLE_DEPTH: usize = FORWARD_DEPTH + 1;
|
||||
|
||||
/// Default chunk size — 16 MiB matches the ISO-mux sector batch and
|
||||
/// is large enough that per-chunk overhead is amortised; small
|
||||
/// enough that the in-flight memory footprint stays bounded.
|
||||
pub const DEFAULT_CHUNK_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
/// Returned from [`BytePrefetcher::into_channels`]. Owns the
|
||||
/// producer-thread join handle so dropping the shell joins the
|
||||
/// producer.
|
||||
pub struct PrefetchShell {
|
||||
producer: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl Drop for PrefetchShell {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.producer.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Spawned byte prefetcher. Drop joins the producer thread.
|
||||
pub struct BytePrefetcher {
|
||||
rx: Receiver<Batch>,
|
||||
recycle_tx: Sender<Vec<u8>>,
|
||||
producer: Option<JoinHandle<()>>,
|
||||
}
|
||||
|
||||
impl BytePrefetcher {
|
||||
/// Spawn the producer thread. `reader` must be `Send` because it
|
||||
/// moves into the thread. `chunk_bytes` is the size of each
|
||||
/// recycled buffer; pick the natural batch size of the
|
||||
/// downstream demuxer (16 MiB for the BD-TS mux pipeline).
|
||||
pub fn new<R: Read + Send + 'static>(
|
||||
mut reader: R,
|
||||
chunk_bytes: usize,
|
||||
halt: Option<Halt>,
|
||||
) -> Self {
|
||||
let (tx, rx) = bounded::<Batch>(FORWARD_DEPTH);
|
||||
let (recycle_tx, recycle_rx) = bounded::<Vec<u8>>(RECYCLE_DEPTH);
|
||||
|
||||
// Seed the recycle pool. Without these the first
|
||||
// `recycle_rx.recv()` would block forever (no consumer has
|
||||
// returned a buffer yet).
|
||||
for _ in 0..RECYCLE_DEPTH {
|
||||
let _ = recycle_tx.send(vec![0u8; chunk_bytes]);
|
||||
}
|
||||
|
||||
let producer = std::thread::Builder::new()
|
||||
.name("freemkv-byte-prefetch".into())
|
||||
.spawn(move || {
|
||||
loop {
|
||||
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
|
||||
return;
|
||||
}
|
||||
let mut buf = match recycle_rx.recv() {
|
||||
Ok(b) => b,
|
||||
Err(_) => return, // consumer dropped both channels
|
||||
};
|
||||
// Re-expose the full extent (previous iteration
|
||||
// may have truncated after a short read).
|
||||
if buf.len() < chunk_bytes {
|
||||
buf.resize(chunk_bytes, 0);
|
||||
} else {
|
||||
// SAFETY: capacity is at least chunk_bytes
|
||||
// after construction.
|
||||
unsafe { buf.set_len(chunk_bytes) };
|
||||
}
|
||||
// Read up to one full chunk. Short reads are
|
||||
// valid and common — pipe `truncate` so the
|
||||
// consumer sees only the bytes that arrived.
|
||||
let n = match reader.read(&mut buf[..]) {
|
||||
Ok(0) => return, // EOF — drop tx, consumer sees RecvError
|
||||
Ok(n) => n,
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(e));
|
||||
return;
|
||||
}
|
||||
};
|
||||
buf.truncate(n);
|
||||
if tx.send(Ok(buf)).is_err() {
|
||||
return; // consumer dropped
|
||||
}
|
||||
}
|
||||
})
|
||||
.expect("freemkv-byte-prefetch thread spawn failed");
|
||||
|
||||
Self {
|
||||
rx,
|
||||
recycle_tx,
|
||||
producer: Some(producer),
|
||||
}
|
||||
}
|
||||
|
||||
/// Peel off the channels for zero-copy pipeline consumption. The
|
||||
/// caller (typically [`crate::mux::demux_thread::DemuxThread`])
|
||||
/// drains `rx`, runs the demuxer in place on each filled buffer,
|
||||
/// and recycles back through `recycle_tx`.
|
||||
pub fn into_channels(self) -> (Receiver<Batch>, Sender<Vec<u8>>, PrefetchShell) {
|
||||
let mut me = self;
|
||||
let producer = me.producer.take();
|
||||
let rx = me.rx.clone();
|
||||
let recycle = me.recycle_tx.clone();
|
||||
std::mem::forget(me);
|
||||
(rx, recycle, PrefetchShell { producer })
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for BytePrefetcher {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.producer.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
//! Linux read-side platform hooks: sequential-access hint at open +
|
||||
//! periodic page-cache eviction during streaming reads.
|
||||
//!
|
||||
//! ## Why both
|
||||
//!
|
||||
//! `POSIX_FADV_SEQUENTIAL` at open widens the kernel's readahead window
|
||||
//! so each pread aggregates into fewer NFS round-trips. `DONTNEED` on
|
||||
//! the consumed window (called periodically by the caller) drops the
|
||||
//! already-read pages from the page cache so an 85 GB streaming ISO
|
||||
//! read doesn't fill memory and starve concurrent writes (the MKV
|
||||
//! output during mux). Together they mirror the write-side
|
||||
//! WritebackPipeline's policy.
|
||||
//!
|
||||
//! ## History
|
||||
//!
|
||||
//! Pre-Phase-1 (0.20.7 baseline) had both. Phase 1's introduction of
|
||||
//! `FileSectorSource` silently dropped the read-side DONTNEED, and
|
||||
//! 0.21.2's revert of `SEQUENTIAL` (mistakenly attributing a regression
|
||||
//! to it) removed the hint. Net effect: 85 GB of ISO reads pinned in
|
||||
//! the page cache + no readahead widening → mux throughput collapse
|
||||
//! from 18 MB/s historical to 2.7-8 MB/s on 0.21.x. Restored in 0.21.6.
|
||||
|
||||
use std::fs::File;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
pub(super) fn hint_sequential(file: &File, _len_bytes: u64) {
|
||||
// Best-effort: return value ignored. A fadvise failure has no
|
||||
// user-observable consequence.
|
||||
unsafe {
|
||||
libc::posix_fadvise(file.as_raw_fd(), 0, 0, libc::POSIX_FADV_SEQUENTIAL);
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop pages in the half-open byte range `[start, start+len)` from
|
||||
/// the page cache. Called periodically by `read_sectors` to bound the
|
||||
/// read-side page cache pressure.
|
||||
pub(super) fn drop_window(file: &File, start: u64, len: u64) {
|
||||
unsafe {
|
||||
libc::posix_fadvise(
|
||||
file.as_raw_fd(),
|
||||
start as i64,
|
||||
len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Async-prefetch `len` bytes at `offset` into the page cache. The
|
||||
/// kernel `readahead(2)` syscall queues the I/O and returns
|
||||
/// immediately — it does NOT wait for completion. Called right after
|
||||
/// each consumed read so the next batch's I/O overlaps with the
|
||||
/// caller's processing of the current batch (decrypt + demux + mux).
|
||||
///
|
||||
/// Without this hint, with a synchronous demux consumer running at
|
||||
/// ~50 MB/s and a single-spindle disk capable of ~150 MB/s, the disk
|
||||
/// sits idle ~70% of each iteration because kernel readahead alone
|
||||
/// (capped at `/sys/block/<dev>/queue/read_ahead_kb`, default 128 KB)
|
||||
/// can only pre-stage a tiny slice of the next batch. An explicit
|
||||
/// `readahead()` of the same size as the current batch tells the
|
||||
/// kernel to queue the full next-batch read now.
|
||||
pub(super) fn prefetch(file: &File, offset: u64, len: u64) {
|
||||
unsafe {
|
||||
libc::readahead(file.as_raw_fd(), offset as i64, len as usize);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
//! macOS: hint the kernel to prefetch a generous chunk. macOS has no
|
||||
//! direct `POSIX_FADV_SEQUENTIAL` equivalent; the idiomatic hint is
|
||||
//! `fcntl(F_RDADVISE, &radvisory)` describing the byte range you
|
||||
//! intend to read soon. We point it at the whole file (clamped to a
|
||||
//! ceiling so a multi-TB ISO doesn't ask the kernel to prefetch
|
||||
//! everything at once).
|
||||
|
||||
use std::fs::File;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
/// `F_RDADVISE` opcode — not in libc's named constants on all SDKs.
|
||||
const F_RDADVISE: libc::c_int = 44;
|
||||
|
||||
/// Cap on the byte length we pass to `F_RDADVISE`. Asking for a
|
||||
/// multi-GB readahead window is counterproductive — the OS doesn't
|
||||
/// have that much cache to throw at one fd. 64 MiB is generous for
|
||||
/// our use case (sweep, mux) and matches the byte-channel cap so the
|
||||
/// kernel's prefetch ≥ our app-level pipeline depth.
|
||||
const RDADVISE_MAX_BYTES: i64 = 64 * 1024 * 1024;
|
||||
|
||||
/// `radvisory` per `<sys/fcntl.h>`. repr(C) layout is stable.
|
||||
#[repr(C)]
|
||||
struct RadAdvisory {
|
||||
ra_offset: libc::off_t,
|
||||
ra_count: libc::c_int,
|
||||
}
|
||||
|
||||
pub(super) fn hint_sequential(file: &File, len_bytes: u64) {
|
||||
let bytes = (len_bytes as i64).min(RDADVISE_MAX_BYTES);
|
||||
let mut ra = RadAdvisory {
|
||||
ra_offset: 0,
|
||||
ra_count: bytes as libc::c_int,
|
||||
};
|
||||
// Best-effort.
|
||||
unsafe {
|
||||
libc::fcntl(file.as_raw_fd(), F_RDADVISE, &mut ra);
|
||||
}
|
||||
}
|
||||
|
||||
/// macOS has no direct `POSIX_FADV_DONTNEED` equivalent for a byte
|
||||
/// range. `fcntl(F_NOCACHE)` would disable caching globally on the fd
|
||||
/// (too coarse — we want the unread region to still benefit). Best
|
||||
/// approximation: no-op. macOS's unified buffer cache is generally
|
||||
/// less prone to the pin-everything pathology that triggers the
|
||||
/// regression on Linux NFS clients.
|
||||
pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {}
|
||||
|
||||
/// Async-prefetch the byte range `[offset, offset+len)`. macOS uses
|
||||
/// the same `fcntl(F_RDADVISE, &radvisory)` primitive as the open-
|
||||
/// time sequential hint, just targeted at a moving window instead of
|
||||
/// the whole file. The kernel queues I/O for the requested range and
|
||||
/// returns immediately.
|
||||
pub(super) fn prefetch(file: &File, offset: u64, len: u64) {
|
||||
let bytes = (len as i64).min(RDADVISE_MAX_BYTES);
|
||||
let mut ra = RadAdvisory {
|
||||
ra_offset: offset as libc::off_t,
|
||||
ra_count: bytes as libc::c_int,
|
||||
};
|
||||
unsafe {
|
||||
libc::fcntl(file.as_raw_fd(), F_RDADVISE, &mut ra);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,361 @@
|
||||
//! [`FileSectorSource`] — read 2048-byte sectors from an ISO file on
|
||||
//! disk via direct `seek + read_exact` (`pread`-equivalent) calls,
|
||||
//! letting the kernel's own readahead policy manage prefetch.
|
||||
//!
|
||||
//! ## Why no app-level buffer
|
||||
//!
|
||||
//! Pre-0.21.3 this source held a 32 MiB (later 4 MiB) read-ahead
|
||||
//! buffer to amortise per-sector NFS round-trips. Empirically that
|
||||
//! buffer hurt: 32 MiB refills bursted the NFS TCP connection hard
|
||||
//! enough to starve the concurrent writer, and even a 4 MiB window
|
||||
//! gave the kernel less freedom to pipeline reads with writes. Direct
|
||||
//! pread per call lets Linux's readahead widen as it detects the
|
||||
//! sequential pattern, and naturally interleaves with writeback.
|
||||
//!
|
||||
//! ## DONTNEED on the consumed window
|
||||
//!
|
||||
//! Without page-cache eviction an 85 GB streaming ISO read pins the
|
||||
//! entire file in memory, starves the concurrent writer, and collapses
|
||||
//! mux throughput (observed: 2.7 MB/s mux on 0.21.5 vs. 70 MB/s
|
||||
//! isolated NFS reads). Every [`READ_DROP_CHUNK_BYTES`] of consumed
|
||||
//! bytes we call `posix_fadvise(DONTNEED)` over that window, mirroring
|
||||
//! the write-side [`crate::io::writeback::WritebackPipeline`] policy.
|
||||
//!
|
||||
//! ## Platform open hint
|
||||
//!
|
||||
//! On `open()` each platform issues its "sequential access expected"
|
||||
//! hint so OS-level readahead widens. The hint and the DONTNEED call
|
||||
//! live in per-OS sibling modules ([`linux::hint_sequential`] et al.)
|
||||
//! — no inline `#[cfg]` in this file.
|
||||
//!
|
||||
//! ## Read-ahead prefetch
|
||||
//!
|
||||
//! After every consumed read we issue an OS-level prefetch hint for
|
||||
//! the next equivalent-sized window (`platform::prefetch`). The
|
||||
//! kernel queues that I/O asynchronously and returns immediately, so
|
||||
//! the next batch's read overlaps with the caller's processing of
|
||||
//! the current batch (decrypt + demux + mux). Without this the disk
|
||||
//! sits idle ~70% of each iteration because kernel SEQUENTIAL
|
||||
//! readahead alone (capped at `read_ahead_kb`, default 128 KB) is
|
||||
//! far smaller than our 16 MiB app-level batch.
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
mod other;
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux as platform;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos as platform;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
use other as platform;
|
||||
#[cfg(target_os = "windows")]
|
||||
use windows as platform;
|
||||
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::sector::SectorSource;
|
||||
|
||||
const SECTOR_SIZE: usize = 2048;
|
||||
|
||||
/// Bytes-read threshold per `posix_fadvise(DONTNEED)` drop on the
|
||||
/// read side. Mirrors `WRITEBACK_CHUNK_BYTES` so the read-side page
|
||||
/// cache stays bounded the same way the write side does.
|
||||
///
|
||||
/// 32 MiB is the empirically tuned value on the rip1 test bed (single
|
||||
/// 7200rpm HDD via SATA): smaller windows (8 / 16 MiB) shorten the
|
||||
/// kernel-readahead overlap and slow the producer; larger windows
|
||||
/// (64 / 128 MiB) let the page cache pin enough of the ISO to
|
||||
/// pressure concurrent writes. Override via `FREEMKV_READ_DROP_CHUNK_MIB`.
|
||||
const READ_DROP_CHUNK_BYTES_DEFAULT: u64 = 32 * 1024 * 1024;
|
||||
|
||||
fn read_drop_chunk_bytes() -> u64 {
|
||||
std::env::var("FREEMKV_READ_DROP_CHUNK_MIB")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.filter(|&n| n > 0)
|
||||
.map(|n| n * 1024 * 1024)
|
||||
.unwrap_or(READ_DROP_CHUNK_BYTES_DEFAULT)
|
||||
}
|
||||
|
||||
/// SectorSource backed by a file (ISO image). Every `read_sectors`
|
||||
/// call is a direct `seek + read_exact` against the underlying file
|
||||
/// — kernel readahead handles prefetch, and every
|
||||
/// [`READ_DROP_CHUNK_BYTES_DEFAULT`] bytes of consumed data the
|
||||
/// platform's `DONTNEED` hook drops the consumed window from the
|
||||
/// page cache to bound memory pressure.
|
||||
pub struct FileSectorSource {
|
||||
file: File,
|
||||
/// Total file size in sectors. Constant after construction;
|
||||
/// surfaced via [`SectorSource::capacity_sectors`].
|
||||
capacity: u32,
|
||||
/// Bytes read since the last DONTNEED drop. Drives the per-
|
||||
/// [`read_drop_chunk_bytes`] page-cache eviction in read_sectors.
|
||||
bytes_read_since_drop: u64,
|
||||
/// File offset at which the current drop window starts. The next
|
||||
/// DONTNEED drops from `drop_window_start` for
|
||||
/// `bytes_read_since_drop` bytes.
|
||||
drop_window_start: u64,
|
||||
/// Cached drop chunk size (resolved from env once at open).
|
||||
drop_chunk_bytes: u64,
|
||||
}
|
||||
|
||||
impl FileSectorSource {
|
||||
/// Open an existing ISO file for reading. Capacity is derived
|
||||
/// from `metadata().len() / 2048`. Returns
|
||||
/// [`Error::IsoTooLarge`] if the file would exceed the 32-bit
|
||||
/// LBA address space (~8 TB).
|
||||
///
|
||||
/// Issues the platform's "sequential access expected" hint on the
|
||||
/// fd (Linux `posix_fadvise(SEQUENTIAL)`, macOS `fcntl(F_RDADVISE)`,
|
||||
/// Windows TODO stub) so the kernel's readahead widens.
|
||||
pub fn open(path: &Path) -> std::io::Result<Self> {
|
||||
let file = File::open(path)?;
|
||||
let len = file.metadata()?.len();
|
||||
let sectors = len / SECTOR_SIZE as u64;
|
||||
if sectors > u32::MAX as u64 {
|
||||
return Err(Error::IsoTooLarge {
|
||||
path: path.to_string_lossy().into_owned(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let capacity = sectors as u32;
|
||||
|
||||
// Best-effort sequential hint. Ignored on platforms without
|
||||
// an equivalent primitive (or where the API exists but the
|
||||
// FS doesn't honour it).
|
||||
platform::hint_sequential(&file, len);
|
||||
|
||||
Ok(Self {
|
||||
file,
|
||||
capacity,
|
||||
bytes_read_since_drop: 0,
|
||||
drop_window_start: 0,
|
||||
drop_chunk_bytes: read_drop_chunk_bytes(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl SectorSource for FileSectorSource {
|
||||
fn capacity_sectors(&self) -> u32 {
|
||||
self.capacity
|
||||
}
|
||||
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
lba: u32,
|
||||
count: u16,
|
||||
out: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
let count = count as u32;
|
||||
let bytes = count as usize * SECTOR_SIZE;
|
||||
debug_assert!(
|
||||
out.len() >= bytes,
|
||||
"FileSectorSource::read_sectors: out len {} < requested {}",
|
||||
out.len(),
|
||||
bytes
|
||||
);
|
||||
if count == 0 {
|
||||
return Ok(0);
|
||||
}
|
||||
let offset = lba as u64 * SECTOR_SIZE as u64;
|
||||
self.file
|
||||
.seek(SeekFrom::Start(offset))
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
self.file
|
||||
.read_exact(&mut out[..bytes])
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
|
||||
// Queue the next batch's read with the kernel before the
|
||||
// caller starts processing what we just returned. readahead()
|
||||
// is non-blocking — it queues I/O and returns, so the kernel
|
||||
// pulls those pages into cache while the consumer (decrypt +
|
||||
// demux + mux) runs. Next read_sectors call hits a warm cache.
|
||||
platform::prefetch(&self.file, offset + bytes as u64, bytes as u64);
|
||||
|
||||
// Periodic page-cache eviction on the read side. Without
|
||||
// this, an 85 GB streaming ISO read pins the entire file in
|
||||
// the kernel page cache, which starves concurrent writes and
|
||||
// collapses mux throughput. Mirrors the write-side
|
||||
// WritebackPipeline's DONTNEED policy.
|
||||
self.bytes_read_since_drop += bytes as u64;
|
||||
if self.bytes_read_since_drop >= self.drop_chunk_bytes {
|
||||
let drop_start = self.drop_window_start;
|
||||
let drop_len = self.bytes_read_since_drop;
|
||||
platform::drop_window(&self.file, drop_start, drop_len);
|
||||
self.drop_window_start = drop_start + drop_len;
|
||||
self.bytes_read_since_drop = 0;
|
||||
}
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
use tempfile::tempdir;
|
||||
|
||||
/// Build a deterministic ISO of `sectors` sectors where sector `n`
|
||||
/// is filled with the byte pattern `((n & 0xff) as u8)`. Lets us
|
||||
/// verify any sector by content alone.
|
||||
fn make_iso(path: &std::path::Path, sectors: u32) {
|
||||
let mut f = std::fs::File::create(path).unwrap();
|
||||
let mut chunk = vec![0u8; SECTOR_SIZE];
|
||||
for n in 0..sectors {
|
||||
let b = (n & 0xff) as u8;
|
||||
chunk.iter_mut().for_each(|c| *c = b);
|
||||
f.write_all(&chunk).unwrap();
|
||||
}
|
||||
f.flush().unwrap();
|
||||
}
|
||||
|
||||
/// Sectors used by spanning-boundary tests. Pick something that
|
||||
/// exercises multi-megabyte reads without making test ISOs huge.
|
||||
/// 8192 sectors = 16 MiB — large enough to cross any readahead
|
||||
/// chunk size we set the kernel hint to.
|
||||
const TEST_SPAN_SECTORS: u32 = 8192;
|
||||
|
||||
#[test]
|
||||
fn sequential_reads_match_file() {
|
||||
let total = TEST_SPAN_SECTORS * 2 + 17;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("seq.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
assert_eq!(src.capacity_sectors(), total);
|
||||
|
||||
let mut got = vec![0u8; SECTOR_SIZE];
|
||||
for lba in 0..total {
|
||||
src.read_sectors(lba, 1, &mut got, false).unwrap();
|
||||
let expected = (lba & 0xff) as u8;
|
||||
assert!(
|
||||
got.iter().all(|b| *b == expected),
|
||||
"sector {lba} content mismatch: expected 0x{expected:02x}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn multi_sector_read_across_chunk_boundary() {
|
||||
let total = TEST_SPAN_SECTORS * 2;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("span.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
|
||||
let span_lba = TEST_SPAN_SECTORS - 2;
|
||||
let mut buf4 = vec![0u8; SECTOR_SIZE * 4];
|
||||
src.read_sectors(span_lba, 4, &mut buf4, false).unwrap();
|
||||
for i in 0..4 {
|
||||
let lba = span_lba + i as u32;
|
||||
let expected = (lba & 0xff) as u8;
|
||||
for b in &buf4[i * SECTOR_SIZE..(i + 1) * SECTOR_SIZE] {
|
||||
assert_eq!(*b, expected, "byte mismatch at sub-sector {i}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn backward_seek_reads_correct_bytes() {
|
||||
// Read forward then jump back: the SectorSource contract is
|
||||
// byte-correctness regardless of access pattern.
|
||||
let total = TEST_SPAN_SECTORS * 2 + 5;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("back.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
let mut got = vec![0u8; SECTOR_SIZE];
|
||||
|
||||
src.read_sectors(TEST_SPAN_SECTORS + 1, 1, &mut got, false)
|
||||
.unwrap();
|
||||
src.read_sectors(0, 1, &mut got, false).unwrap();
|
||||
assert!(got.iter().all(|b| *b == 0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_at_eof_returns_correct_bytes() {
|
||||
// File smaller than the readahead chunk — reads near EOF must
|
||||
// still return correct bytes.
|
||||
let total: u32 = 100;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("small.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
assert_eq!(src.capacity_sectors(), total);
|
||||
|
||||
let mut got = vec![0u8; SECTOR_SIZE];
|
||||
src.read_sectors(0, 1, &mut got, false).unwrap();
|
||||
src.read_sectors(total - 1, 1, &mut got, false).unwrap();
|
||||
let expected = ((total - 1) & 0xff) as u8;
|
||||
assert!(got.iter().all(|b| *b == expected));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn large_single_read() {
|
||||
// A multi-MB single read must work — the implementation has
|
||||
// no app-level chunking, so this just exercises the direct
|
||||
// pread path on a larger request.
|
||||
let total = TEST_SPAN_SECTORS + 100;
|
||||
let dir = tempdir().unwrap();
|
||||
let path = dir.path().join("big.iso");
|
||||
make_iso(&path, total);
|
||||
|
||||
let mut src = FileSectorSource::open(&path).unwrap();
|
||||
let req = (TEST_SPAN_SECTORS + 1) as u16;
|
||||
let req_bytes = req as usize * SECTOR_SIZE;
|
||||
let mut big = vec![0u8; req_bytes];
|
||||
src.read_sectors(0, req, &mut big, false).unwrap();
|
||||
assert!(big[..SECTOR_SIZE].iter().all(|b| *b == 0));
|
||||
let last_lba = req as u32 - 1;
|
||||
let exp = (last_lba & 0xff) as u8;
|
||||
let last_off = (req as usize - 1) * SECTOR_SIZE;
|
||||
assert!(
|
||||
big[last_off..last_off + SECTOR_SIZE]
|
||||
.iter()
|
||||
.all(|b| *b == exp)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_chunk_size_env_override() {
|
||||
// Explicit 8 MiB via env var.
|
||||
// SAFETY: tests in this crate are single-threaded per the
|
||||
// default cargo test harness, but std::env::set_var is
|
||||
// declared `unsafe` since Rust 2024 (it can race with other
|
||||
// threads / TLS). For a test that runs in-process before any
|
||||
// FileSectorSource construction this is safe in practice.
|
||||
unsafe {
|
||||
std::env::set_var("FREEMKV_READ_DROP_CHUNK_MIB", "8");
|
||||
}
|
||||
assert_eq!(read_drop_chunk_bytes(), 8 * 1024 * 1024);
|
||||
|
||||
unsafe {
|
||||
std::env::remove_var("FREEMKV_READ_DROP_CHUNK_MIB");
|
||||
}
|
||||
assert_eq!(read_drop_chunk_bytes(), READ_DROP_CHUNK_BYTES_DEFAULT);
|
||||
|
||||
// Garbage env value falls back to default.
|
||||
unsafe {
|
||||
std::env::set_var("FREEMKV_READ_DROP_CHUNK_MIB", "not-a-number");
|
||||
}
|
||||
assert_eq!(read_drop_chunk_bytes(), READ_DROP_CHUNK_BYTES_DEFAULT);
|
||||
unsafe {
|
||||
std::env::remove_var("FREEMKV_READ_DROP_CHUNK_MIB");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
//! Fallback for targets without a known sequential-readahead hint
|
||||
//! (BSDs, illumos, etc.). No-op — reads still work, they just don't
|
||||
//! get the OS-level prefetch widening.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
pub(super) fn hint_sequential(_file: &File, _len_bytes: u64) {}
|
||||
|
||||
pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {}
|
||||
|
||||
pub(super) fn prefetch(_file: &File, _offset: u64, _len: u64) {}
|
||||
@@ -0,0 +1,30 @@
|
||||
//! Windows: the canonical sequential-access hint is
|
||||
//! `FILE_FLAG_SEQUENTIAL_SCAN` passed to `CreateFile` at open time —
|
||||
//! it cannot be set after the fact via `SetFileInformationByHandle`.
|
||||
//! Routing the open call through this module would mean a custom
|
||||
//! `File::from_raw_handle` plumb for every `FileSectorSource::open`
|
||||
//! caller, which is more invasive than the Phase 1 scope.
|
||||
//!
|
||||
//! TODO: replumb `FileSectorSource::open` to take an
|
||||
//! `OpenOptions`-style builder so the Windows path can flip the flag
|
||||
//! at open time. For now this is a no-op stub.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
pub(super) fn hint_sequential(_file: &File, _len_bytes: u64) {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"FileSectorSource hint_sequential: windows stub (TODO: FILE_FLAG_SEQUENTIAL_SCAN at open)"
|
||||
);
|
||||
}
|
||||
|
||||
/// Windows page-cache eviction is not exposed via a posix_fadvise
|
||||
/// equivalent. The kernel does its own working-set management. No-op
|
||||
/// for now.
|
||||
pub(super) fn drop_window(_file: &File, _start: u64, _len: u64) {}
|
||||
|
||||
/// Windows async-prefetch hint. With FILE_FLAG_SEQUENTIAL_SCAN at
|
||||
/// open the kernel already prefetches aggressively, so there's no
|
||||
/// per-range hint we'd add on top. No-op stub for parity with the
|
||||
/// posix platforms.
|
||||
pub(super) fn prefetch(_file: &File, _offset: u64, _len: u64) {}
|
||||
+21
-6
@@ -10,24 +10,39 @@
|
||||
//! no-op pipeline — their default cache policies have not been shown
|
||||
//! to exhibit the same pathology for this access pattern.
|
||||
//!
|
||||
//! `FileSectorSource` is the read-side dual — it implements
|
||||
//! [`crate::sector::SectorSource`] for an ISO file using direct
|
||||
//! `pread`-equivalent calls so the kernel's own readahead policy runs
|
||||
//! (which interleaves naturally with the concurrent writeback). It
|
||||
//! pairs that with periodic `posix_fadvise(DONTNEED)` drops on the
|
||||
//! consumed window so an 85 GB streaming ISO read doesn't fill the
|
||||
//! page cache and starve the concurrent MKV write.
|
||||
//!
|
||||
//! `Pipeline` + `Sink` (0.18) is the generic producer/consumer primitive
|
||||
//! used by sweep, patch, and mux to overlap reads with writes via a
|
||||
//! bounded channel + dedicated consumer thread.
|
||||
//!
|
||||
//! `byte_channel` is a byte-sized producer/consumer channel for the
|
||||
//! mux pipeline, sized to absorb worst-case input read stalls (see
|
||||
//! `(internal)/memory/project_buffering_architecture.md`).
|
||||
|
||||
pub(crate) mod bounded;
|
||||
pub mod byte_channel;
|
||||
pub mod byte_prefetcher;
|
||||
pub mod file_sector_source;
|
||||
pub mod sink;
|
||||
mod writeback;
|
||||
mod writeback_file;
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
pub(crate) mod platform_macos;
|
||||
|
||||
pub mod pipeline;
|
||||
|
||||
pub(crate) use writeback_file::WritebackFile;
|
||||
|
||||
// Re-exports for the 0.18 redesign. Sweep + patch are both wired up
|
||||
// (disc/sweep.rs, disc/patch.rs); mux migrates separately in autorip.
|
||||
// `WRITE_THROUGH_DEPTH` is patch-specific and has no other in-tree
|
||||
// caller — the targeted `#[allow]` keeps the re-export visible without
|
||||
// dragging the rest of the module under `dead_code`.
|
||||
#[allow(unused_imports)]
|
||||
// Re-exports for the 0.18 redesign. Sweep, patch, and mux are all
|
||||
// wired up (disc/sweep.rs, disc/patch.rs, autorip's ripper/mux.rs).
|
||||
pub use pipeline::{
|
||||
DEFAULT_PIPELINE_DEPTH, Flow, Pipeline, READ_PIPELINE_DEPTH, Sink, WRITE_PIPELINE_DEPTH,
|
||||
WRITE_THROUGH_DEPTH,
|
||||
|
||||
+72
-49
@@ -34,10 +34,11 @@
|
||||
//! consumer lag detection). This is critical for diagnosing stalls.
|
||||
|
||||
use std::io;
|
||||
use std::sync::mpsc::{SyncSender, TrySendError, sync_channel};
|
||||
use std::thread::{self, JoinHandle};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crossbeam_channel::{Sender, TrySendError, bounded};
|
||||
|
||||
use crate::error::Error;
|
||||
use crate::halt::Halt;
|
||||
|
||||
@@ -52,17 +53,26 @@ use crate::halt::Halt;
|
||||
/// caller has already lost the rip.
|
||||
pub const JOIN_TIMEOUT_SECS: u64 = 600;
|
||||
|
||||
/// Polling slice for the halt-aware send/finish loops. Mirrors the
|
||||
/// `bounded_syscall` cadence (250 ms) so halt observation feels equally
|
||||
/// responsive across both primitives.
|
||||
const POLL_INTERVAL: Duration = Duration::from_millis(250);
|
||||
|
||||
/// Polling slice for the halt-aware send loop. Smaller than
|
||||
/// [`POLL_INTERVAL`] because send latency tolerance is much lower —
|
||||
/// frames must move through the channel at hundreds-of-Hz on the
|
||||
/// happy path; 50 ms keeps backpressure-driven wakeups fine-grained
|
||||
/// without busy-looping.
|
||||
const SEND_POLL_INTERVAL: Duration = Duration::from_millis(50);
|
||||
/// Halt-check cadence for the send loop. Producer blocks on
|
||||
/// [`crossbeam_channel::Sender::send_timeout`] for this slice — the
|
||||
/// kernel wakes it the instant the consumer drains a slot, so on the
|
||||
/// happy path there's no throughput cap from this primitive at all
|
||||
/// (the cap is whatever the underlying medium can sustain). When the
|
||||
/// consumer is genuinely wedged, the timeout fires every
|
||||
/// [`crate::halt::POLL_INTERVAL`] and the producer checks the halt
|
||||
/// token; that's the latency a stop request will observe.
|
||||
///
|
||||
/// Single source of truth lives in [`crate::halt::POLL_INTERVAL`]
|
||||
/// (also used by `bounded_syscall`). Aliased here for readability of
|
||||
/// the send/finish call sites below.
|
||||
///
|
||||
/// 0.21.7 replaced an old `std::sync::mpsc::sync_channel` + 50 ms
|
||||
/// `thread::sleep` polling loop that capped mux throughput at
|
||||
/// ~20 frames/sec ≈ 1 MB/s on saturated channels. See
|
||||
/// (internal)/memory/feedback_send_with_halt_poll_throttle.md
|
||||
/// for the multi-day diagnostic that surfaced it.
|
||||
use crate::halt::POLL_INTERVAL;
|
||||
const SEND_HALT_CHECK_INTERVAL: Duration = POLL_INTERVAL;
|
||||
|
||||
/// Check if verbose debug logging is enabled via FREEMKV_DEBUG env var.
|
||||
pub fn debug_enabled() -> bool {
|
||||
@@ -90,8 +100,7 @@ pub const WRITE_PIPELINE_DEPTH: usize = 16;
|
||||
/// Channel depth for write-through pipelines. Each `send` fully
|
||||
/// drains before the next can enqueue. Use this when the producer
|
||||
/// must observe consumer side-effects (e.g. mapfile state) before
|
||||
/// emitting the next item.
|
||||
#[allow(dead_code)]
|
||||
/// emitting the next item. Currently used by `disc::patch`.
|
||||
pub const WRITE_THROUGH_DEPTH: usize = 1;
|
||||
|
||||
/// Outcome of [`Sink::apply`]: either keep feeding items
|
||||
@@ -135,7 +144,7 @@ pub trait Sink<I>: Send + 'static {
|
||||
/// Bounded producer/consumer pipeline. Holds the producer-side
|
||||
/// channel and the consumer thread's join handle.
|
||||
pub struct Pipeline<I: Send + 'static, R: Send + 'static> {
|
||||
tx: SyncSender<I>,
|
||||
tx: Sender<I>,
|
||||
handle: JoinHandle<Result<R, Error>>,
|
||||
}
|
||||
|
||||
@@ -152,11 +161,10 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
/// propagated rather than panicked.
|
||||
///
|
||||
/// Sweep uses [`Pipeline::spawn_named`] directly so the consumer
|
||||
/// thread shows up as `freemkv-sweep-consumer`; this function has
|
||||
/// no in-tree caller yet. Patch and mux migrate in later 0.18
|
||||
/// slices. The targeted `#[allow]` is removed when one of them
|
||||
/// lands on the default name.
|
||||
#[allow(dead_code)]
|
||||
/// thread shows up as `freemkv-sweep-consumer`; mux uses
|
||||
/// `freemkv-mux-consumer`. `Pipeline::spawn` (this function, with
|
||||
/// the default name) is used by `disc::patch` and by the unit
|
||||
/// tests in this module.
|
||||
pub fn spawn<S: Sink<I, Output = R>>(depth: usize, sink: S) -> Result<Self, Error> {
|
||||
Self::spawn_named("freemkv-pipeline-consumer", depth, sink)
|
||||
}
|
||||
@@ -170,7 +178,7 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
depth: usize,
|
||||
sink: S,
|
||||
) -> Result<Self, Error> {
|
||||
let (tx, rx) = sync_channel::<I>(depth);
|
||||
let (tx, rx) = bounded::<I>(depth);
|
||||
let handle = thread::Builder::new()
|
||||
.name(name.into())
|
||||
.spawn(move || -> Result<R, Error> {
|
||||
@@ -281,14 +289,20 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
/// `Err`. Useful for best-effort signalling (e.g. sweep's
|
||||
/// throttled `StatsRequest`) where dropping the message is
|
||||
/// preferable to blocking the producer.
|
||||
pub fn try_send(&self, item: I) -> Result<(), std::sync::mpsc::TrySendError<I>> {
|
||||
pub fn try_send(&self, item: I) -> Result<(), TrySendError<I>> {
|
||||
self.tx.try_send(item)
|
||||
}
|
||||
|
||||
/// Halt-aware bounded variant of [`Pipeline::send`].
|
||||
///
|
||||
/// Polls `try_send` on a 50 ms slice. Between slices, checks
|
||||
/// (1) the [`Halt`] token and (2) the per-call `deadline`. Returns:
|
||||
/// Uses [`crossbeam_channel::Sender::send_timeout`] so the producer
|
||||
/// thread BLOCKS on consumer drain (kernel-wakeup) rather than
|
||||
/// polling. The timeout slice is just the halt-observation cadence
|
||||
/// ([`SEND_HALT_CHECK_INTERVAL`]) — on the happy path the producer
|
||||
/// wakes the instant the consumer drains a slot, so there is no
|
||||
/// throughput cap from this primitive at any medium speed.
|
||||
///
|
||||
/// Returns:
|
||||
///
|
||||
/// - `Ok(())` once the item lands in the channel.
|
||||
/// - `Err(item)` if the consumer disconnected, the halt fired, or
|
||||
@@ -302,36 +316,45 @@ impl<I: Send + 'static, R: Send + 'static> Pipeline<I, R> {
|
||||
/// Unlike [`Pipeline::send`], this never blocks the producer
|
||||
/// thread inside an unkillable `mpsc::send` — if the consumer is
|
||||
/// wedged inside an unkillable syscall, the producer can still
|
||||
/// observe `/api/stop` and unwind.
|
||||
/// observe `/api/stop` and unwind within
|
||||
/// [`SEND_HALT_CHECK_INTERVAL`].
|
||||
pub fn send_with_halt(&self, item: I, halt: &Halt, deadline: Duration) -> Result<(), I> {
|
||||
use crossbeam_channel::SendTimeoutError;
|
||||
let end = Instant::now() + deadline;
|
||||
let mut pending = item;
|
||||
loop {
|
||||
match self.tx.try_send(pending) {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(TrySendError::Full(returned)) => {
|
||||
pending = returned;
|
||||
if halt.is_cancelled() {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: halt observed, returning item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
return Err(pending);
|
||||
}
|
||||
if Instant::now() >= end {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: deadline elapsed, returning item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
return Err(pending);
|
||||
}
|
||||
thread::sleep(SEND_POLL_INTERVAL);
|
||||
// Pre-check the cheap exit conditions before parking.
|
||||
if halt.is_cancelled() {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: halt observed, returning item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
Err(TrySendError::Disconnected(returned)) => {
|
||||
return Err(pending);
|
||||
}
|
||||
let now = Instant::now();
|
||||
if now >= end {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: deadline elapsed, returning item={}",
|
||||
std::any::type_name::<I>()
|
||||
);
|
||||
}
|
||||
return Err(pending);
|
||||
}
|
||||
// Wait for space-available or halt-check tick, whichever
|
||||
// is sooner. Crossbeam's send_timeout is kernel-wakeup
|
||||
// based: the consumer's recv on a saturated channel
|
||||
// signals this thread the moment a slot opens up.
|
||||
let slice = SEND_HALT_CHECK_INTERVAL.min(end.saturating_duration_since(now));
|
||||
match self.tx.send_timeout(pending, slice) {
|
||||
Ok(()) => return Ok(()),
|
||||
Err(SendTimeoutError::Timeout(returned)) => {
|
||||
pending = returned;
|
||||
// loop: re-check halt + deadline, then park again
|
||||
}
|
||||
Err(SendTimeoutError::Disconnected(returned)) => {
|
||||
if debug_enabled() {
|
||||
tracing::debug!(
|
||||
"Pipeline send_with_halt: consumer disconnected, item={}",
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
//! Shared macOS `fcntl(F_PREALLOCATE)` definitions.
|
||||
//!
|
||||
//! The `libc` crate doesn't expose these symbols across all macOS SDK
|
||||
//! versions, so we define them locally with values from
|
||||
//! `/usr/include/sys/fcntl.h`. Two call sites (
|
||||
//! [`crate::io::writeback_file`] and [`crate::io::sink::preallocate`])
|
||||
//! need the same constants and `fstore_t` layout — keeping a single
|
||||
//! source of truth here prevents the two copies from drifting.
|
||||
//!
|
||||
//! Module-level cfg gate lives in the parent (`io/mod.rs`); this file
|
||||
//! is only compiled on macOS, so no inner `#![cfg]` is needed.
|
||||
|
||||
/// `fcntl(F_PREALLOCATE)` command number from `sys/fcntl.h`.
|
||||
pub(crate) const F_PREALLOCATE: libc::c_int = 42;
|
||||
|
||||
/// Anchor preallocation at the current physical EOF.
|
||||
pub(crate) const F_PEOFPOSMODE: libc::c_int = 3;
|
||||
|
||||
/// Prefer a contiguous allocation. Try this first; on `EINVAL` (no
|
||||
/// contiguous run of that size), fall back to `F_ALLOCATEALL`.
|
||||
pub(crate) const F_ALLOCATECONTIG: libc::c_uint = 0x0000_0002;
|
||||
|
||||
/// Allow non-contiguous allocation. Stronger guarantee than just
|
||||
/// asking for `F_ALLOCATECONTIG` because the kernel will piece
|
||||
/// together fragments rather than failing.
|
||||
pub(crate) const F_ALLOCATEALL: libc::c_uint = 0x0000_0004;
|
||||
|
||||
/// `fstore_t` from `sys/fcntl.h`. `repr(C)` because we hand it to
|
||||
/// `fcntl(F_PREALLOCATE)` which writes through the pointer.
|
||||
#[repr(C)]
|
||||
#[derive(Clone, Copy)]
|
||||
pub(crate) struct Fstore {
|
||||
pub fst_flags: libc::c_uint,
|
||||
pub fst_posmode: libc::c_int,
|
||||
pub fst_offset: libc::off_t,
|
||||
pub fst_length: libc::off_t,
|
||||
pub fst_bytesalloc: libc::off_t,
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
//! `LocalFileSink` — `BufWriter<File>` for the common local-disk case.
|
||||
//!
|
||||
//! Buffering: 4 MiB internal `BufWriter`. Sized to coalesce the small
|
||||
//! per-PES writes that come out of the muxer into kernel-page-aligned
|
||||
//! flushes without making the buffer big enough to matter for memory
|
||||
//! pressure on a single concurrent rip.
|
||||
//!
|
||||
//! `Seek` flushes the underlying `BufWriter` first; otherwise a seek
|
||||
//! could leapfrog buffered data and silently corrupt the file. This is
|
||||
//! the same shape `BufWriter` itself uses when it impls `Seek` in
|
||||
//! stdlib, and is necessary for MKV's seek-back operations (cluster
|
||||
//! size patch, Cues index, segment header backpatch) to land on the
|
||||
//! right offset.
|
||||
//!
|
||||
//! `RandomAccessSink` is satisfied via the blanket impl in
|
||||
//! [`super::mod`]; no explicit impl needed here.
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{self, BufWriter, Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use super::preallocate;
|
||||
|
||||
const BUFFER_BYTES: usize = 4 * 1024 * 1024;
|
||||
|
||||
/// Random-access write sink for local disks.
|
||||
///
|
||||
/// Wraps a `BufWriter<File>` with a 4 MiB internal buffer and forwards
|
||||
/// `Write`/`Seek` so any call site that previously held a `File` or
|
||||
/// `WritebackFile` can drop this in. `finish()` flushes the buffer and
|
||||
/// runs `sync_all` on the underlying file so the caller can drop it
|
||||
/// without losing data.
|
||||
///
|
||||
/// Construction always opens the file `create + truncate + read +
|
||||
/// write`. `read` is enabled so the same handle can be reused for a
|
||||
/// verification re-read after the mux (the existing
|
||||
/// `FileSectorSink::create` pattern). On Linux, [`with_size_hint`]
|
||||
/// additionally calls `fallocate(FALLOC_FL_KEEP_SIZE)` to pre-reserve
|
||||
/// extents.
|
||||
///
|
||||
/// [`with_size_hint`]: Self::with_size_hint
|
||||
pub struct LocalFileSink {
|
||||
inner: BufWriter<File>,
|
||||
}
|
||||
|
||||
impl LocalFileSink {
|
||||
/// Open `path` for writing, truncating any existing contents.
|
||||
pub fn create(path: &Path) -> io::Result<Self> {
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
Ok(Self {
|
||||
inner: BufWriter::with_capacity(BUFFER_BYTES, file),
|
||||
})
|
||||
}
|
||||
|
||||
/// Like [`Self::create`] but additionally calls the per-OS
|
||||
/// preallocate path with `size_bytes`. On Linux this is
|
||||
/// `fallocate(FALLOC_FL_KEEP_SIZE)` so the on-disk extents are
|
||||
/// reserved up front (reducing fragmentation for big sequential
|
||||
/// muxer output); on other OSes it is a no-op today. Failures
|
||||
/// from the preallocate call are non-fatal — the file is still
|
||||
/// returned, just without the size reservation.
|
||||
pub fn with_size_hint(path: &Path, size_bytes: u64) -> io::Result<Self> {
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.truncate(true)
|
||||
.open(path)?;
|
||||
preallocate::preallocate(&file, size_bytes);
|
||||
Ok(Self {
|
||||
inner: BufWriter::with_capacity(BUFFER_BYTES, file),
|
||||
})
|
||||
}
|
||||
|
||||
/// Drain the internal buffer and `fsync` the underlying file.
|
||||
/// Idempotent with `Drop` (the `BufWriter` also flushes on drop;
|
||||
/// this call additionally surfaces fsync errors to the caller).
|
||||
#[allow(dead_code)] // exposed for parity with WritebackFile::sync_all
|
||||
pub fn sync_all(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()?;
|
||||
self.inner.get_ref().sync_all()
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for LocalFileSink {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.inner.write(buf)
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.inner.write_all(buf)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl Seek for LocalFileSink {
|
||||
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
|
||||
// Flush before seeking so buffered bytes land at the offset
|
||||
// they were written for, not the new one.
|
||||
self.inner.flush()?;
|
||||
self.inner.get_mut().seek(from)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Read;
|
||||
|
||||
#[test]
|
||||
fn write_seek_roundtrip() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("rt.bin");
|
||||
let mut s = LocalFileSink::create(&p).unwrap();
|
||||
s.write_all(b"AAAA").unwrap();
|
||||
s.write_all(b"BBBB").unwrap();
|
||||
// Seek back over the second word and overwrite.
|
||||
s.seek(SeekFrom::Start(4)).unwrap();
|
||||
s.write_all(b"CCCC").unwrap();
|
||||
s.sync_all().unwrap();
|
||||
drop(s);
|
||||
|
||||
let mut f = File::open(&p).unwrap();
|
||||
let mut got = Vec::new();
|
||||
f.read_to_end(&mut got).unwrap();
|
||||
assert_eq!(&got[..], b"AAAACCCC");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn drop_flushes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("drop.bin");
|
||||
{
|
||||
let mut s = LocalFileSink::create(&p).unwrap();
|
||||
s.write_all(b"buffered").unwrap();
|
||||
// No explicit flush / sync_all — BufWriter drop runs the
|
||||
// flush and the file should land on disk.
|
||||
}
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
assert_eq!(&bytes[..], b"buffered");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_size_hint_creates_writable_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("sz.bin");
|
||||
let mut s = LocalFileSink::with_size_hint(&p, 64 * 1024).unwrap();
|
||||
s.write_all(b"hint-ok").unwrap();
|
||||
s.sync_all().unwrap();
|
||||
drop(s);
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
assert_eq!(&bytes[..], b"hint-ok");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
//! Output-sink trait split for the buffering architecture.
|
||||
//!
|
||||
//! Two traits, one for each capability axis of an output destination:
|
||||
//!
|
||||
//! - [`SequentialSink`] — anything you can `Write` to in order. Sockets,
|
||||
//! pipes, append-only stores, plain files. Containers that don't need
|
||||
//! seek (M2TS, fMP4, HEVC elementary) target this.
|
||||
//! - [`RandomAccessSink`] — everything `SequentialSink` plus a working
|
||||
//! `Seek`. Local files, NFS files, anything with random-write
|
||||
//! semantics. Containers that need backpatch (MKV cluster sizes, Cues
|
||||
//! index, MP4 moov-at-end) target this.
|
||||
//!
|
||||
//! `RandomAccessSink: SequentialSink` — every random-access sink is
|
||||
//! also a valid sequential sink. The muxer is generic over which it
|
||||
//! requires (`MkvMux<S: RandomAccessSink>`, `M2tsMux<S: SequentialSink>`)
|
||||
//! so an attempt to mux MKV to a network socket is a compile error.
|
||||
//!
|
||||
//! Buffering policy belongs to the concrete sink, not to a wrapper at
|
||||
//! the call site. `LocalFileSink` wraps a `BufWriter<File>` with a
|
||||
//! 4 MiB buffer for the common local-disk case; `WritebackFile`
|
||||
//! (separate module) wraps a `File` with the adaptive-chunk
|
||||
//! `sync_file_range` machinery for the Linux+NFS case.
|
||||
//!
|
||||
//! See `(internal)/memory/project_buffering_architecture.md` for
|
||||
//! the full design and the source/sink matrix.
|
||||
|
||||
use std::io::{Seek, Write};
|
||||
|
||||
mod local_file;
|
||||
mod preallocate;
|
||||
mod socket;
|
||||
|
||||
pub use local_file::LocalFileSink;
|
||||
pub use socket::{SocketSink, UdpSocketSink};
|
||||
|
||||
/// Sequential-only write destination. Sockets, pipes, append-only
|
||||
/// stores. No seek. Implementations own their write buffering — the
|
||||
/// trait does not impose or hide any buffering of its own.
|
||||
///
|
||||
/// `finish` drains any internal buffering and signals end-of-stream to
|
||||
/// the underlying transport (close-write on a socket, flush on a
|
||||
/// buffered writer, etc.). The default impl is a no-op; concrete
|
||||
/// implementations that need explicit shutdown can override it but the
|
||||
/// blanket impl below keeps it optional for adapter types like
|
||||
/// `&mut File`.
|
||||
pub trait SequentialSink: Write + Send {
|
||||
fn finish(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Random-access write destination. Local files, NFS files, anything
|
||||
/// with a working `Seek`. Inherits the `SequentialSink` contract — a
|
||||
/// random-access sink is always usable as a sequential sink.
|
||||
pub trait RandomAccessSink: SequentialSink + Seek {}
|
||||
|
||||
// Blanket impls so any `Write + Send` type acts as a `SequentialSink`
|
||||
// (with default `finish`), and any sink that also impls `Seek` is
|
||||
// automatically a `RandomAccessSink`. Keeps call-site ergonomics simple
|
||||
// — `&mut File`, `LocalFileSink`, `WritebackFile`, `BufWriter<File>`,
|
||||
// and `Cursor<Vec<u8>>` all satisfy the right trait without per-type
|
||||
// boilerplate.
|
||||
impl<T: Write + Send + ?Sized> SequentialSink for T {}
|
||||
impl<T: SequentialSink + Seek + ?Sized> RandomAccessSink for T {}
|
||||
|
||||
/// Pick the right `RandomAccessSink` impl for `dest` based on its
|
||||
/// filesystem type.
|
||||
///
|
||||
/// - Linux + NFS path → `WritebackFile` with its adaptive-chunk
|
||||
/// sync_file_range machinery and (when supported) `fallocate` size
|
||||
/// hint.
|
||||
/// - everything else → [`LocalFileSink`] over `BufWriter<File>`. On
|
||||
/// non-Linux there is no `WritebackFile` machinery to opt into, and
|
||||
/// on local Linux the kernel's default writeback policy is already
|
||||
/// fine.
|
||||
///
|
||||
/// `size_hint`, when present, is forwarded to the per-OS preallocate
|
||||
/// path (`fallocate(KEEP_SIZE)` on Linux, `F_PREALLOCATE` on macOS when
|
||||
/// implemented, no-op elsewhere).
|
||||
///
|
||||
/// Returns a boxed trait object so the call site (mux construction)
|
||||
/// stays agnostic of which concrete sink got picked.
|
||||
#[allow(dead_code)] // wiring to mux::resolve is a follow-up commit
|
||||
pub fn open_for_mkv(
|
||||
dest: &std::path::Path,
|
||||
size_hint: Option<u64>,
|
||||
) -> std::io::Result<Box<dyn RandomAccessSink>> {
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
use crate::platform::fs_type::detect;
|
||||
#[cfg(target_os = "linux")]
|
||||
use crate::platform::fs_type::{FsType, detect};
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
if detect(dest) == FsType::Nfs {
|
||||
let wf = match size_hint {
|
||||
Some(n) => crate::io::WritebackFile::create_with_size_hint(dest, n)?,
|
||||
None => crate::io::WritebackFile::create(dest)?,
|
||||
};
|
||||
return Ok(Box::new(wf));
|
||||
}
|
||||
}
|
||||
// Silence the unused-binding warning on non-Linux where the only
|
||||
// branch above is cfg-gated out.
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
{
|
||||
let _ = detect(dest);
|
||||
}
|
||||
|
||||
let sink = match size_hint {
|
||||
Some(n) => LocalFileSink::with_size_hint(dest, n)?,
|
||||
None => LocalFileSink::create(dest)?,
|
||||
};
|
||||
Ok(Box::new(sink))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::fs::File;
|
||||
|
||||
// Type-level assertion: the blanket impls cover the shapes we care
|
||||
// about. These functions never run; they just have to type-check.
|
||||
fn _assert_file_is_sequential(_: &mut dyn SequentialSink) {}
|
||||
fn _assert_file_is_random_access(_: &mut dyn RandomAccessSink) {}
|
||||
|
||||
#[test]
|
||||
fn blanket_impls_cover_file_and_localfilesink() {
|
||||
// `File` directly via blanket impls.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let mut f = File::create(dir.path().join("a.bin")).unwrap();
|
||||
_assert_file_is_sequential(&mut f);
|
||||
_assert_file_is_random_access(&mut f);
|
||||
|
||||
// `LocalFileSink` ditto.
|
||||
let mut s = LocalFileSink::create(&dir.path().join("b.bin")).unwrap();
|
||||
_assert_file_is_sequential(&mut s);
|
||||
_assert_file_is_random_access(&mut s);
|
||||
|
||||
// `WritebackFile` — confirms the Phase 1 type still satisfies
|
||||
// the trait via the blanket impl without needing an explicit
|
||||
// `impl RandomAccessSink for WritebackFile {}`.
|
||||
let mut wf = crate::io::WritebackFile::create(&dir.path().join("c.bin")).unwrap();
|
||||
_assert_file_is_sequential(&mut wf);
|
||||
_assert_file_is_random_access(&mut wf);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn open_for_mkv_returns_a_random_access_sink() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("c.bin");
|
||||
let mut sink = open_for_mkv(&p, Some(64 * 1024)).unwrap();
|
||||
use std::io::{Seek, SeekFrom, Write};
|
||||
sink.write_all(b"hello").unwrap();
|
||||
sink.seek(SeekFrom::Start(0)).unwrap();
|
||||
sink.finish().unwrap();
|
||||
drop(sink);
|
||||
let bytes = std::fs::read(&p).unwrap();
|
||||
assert_eq!(&bytes[..5], b"hello");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//! Linux `fallocate(FALLOC_FL_KEEP_SIZE)` preallocation.
|
||||
//!
|
||||
//! `KEEP_SIZE` reserves extents without changing the apparent file
|
||||
//! length, which matches the muxer's expectation that writes still grow
|
||||
//! the file naturally.
|
||||
|
||||
use std::fs::File;
|
||||
#[cfg(unix)]
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
pub(super) fn preallocate_impl(file: &File, size_bytes: u64) {
|
||||
let fd = file.as_raw_fd();
|
||||
// FALLOC_FL_KEEP_SIZE = 0x01.
|
||||
let rc = unsafe { libc::fallocate(fd, libc::FALLOC_FL_KEEP_SIZE, 0, size_bytes as i64) };
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"LocalFileSink fallocate size_hint={size_bytes} rc={rc} ok={}",
|
||||
rc == 0
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
//! macOS `F_PREALLOCATE` extent reservation.
|
||||
//!
|
||||
//! `fcntl(F_PREALLOCATE)` with `F_ALLOCATECONTIG` first (try for a
|
||||
//! contiguous run) and fall back to `F_ALLOCATEALL` (non-contig OK).
|
||||
//! Reported file size is unchanged — the muxer's writes still grow it.
|
||||
|
||||
use std::fs::File;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
|
||||
use crate::io::platform_macos::{
|
||||
F_ALLOCATEALL, F_ALLOCATECONTIG, F_PEOFPOSMODE, F_PREALLOCATE, Fstore,
|
||||
};
|
||||
|
||||
pub(super) fn preallocate_impl(file: &File, size_bytes: u64) {
|
||||
let fd = file.as_raw_fd();
|
||||
let mut store = Fstore {
|
||||
fst_flags: F_ALLOCATECONTIG,
|
||||
fst_posmode: F_PEOFPOSMODE,
|
||||
fst_offset: 0,
|
||||
fst_length: size_bytes as libc::off_t,
|
||||
fst_bytesalloc: 0,
|
||||
};
|
||||
let mut rc = unsafe { libc::fcntl(fd, F_PREALLOCATE, &mut store as *mut Fstore) };
|
||||
if rc == -1 {
|
||||
// Fall back to non-contiguous.
|
||||
store.fst_flags = F_ALLOCATEALL;
|
||||
rc = unsafe { libc::fcntl(fd, F_PREALLOCATE, &mut store as *mut Fstore) };
|
||||
}
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"LocalFileSink F_PREALLOCATE size_hint={size_bytes} rc={rc} bytesalloc={}",
|
||||
store.fst_bytesalloc
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
//! Per-OS extent preallocation. Best-effort; failures are logged at
|
||||
//! debug and otherwise swallowed because the file is still usable
|
||||
//! without the size reservation — only large-file fragmentation gets
|
||||
//! marginally worse.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
mod other;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux::preallocate_impl;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos::preallocate_impl;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
use other::preallocate_impl;
|
||||
|
||||
/// Reserve `size_bytes` of disk space for `file`'s on-disk extents.
|
||||
/// Reported file size is unchanged — writes still grow the file
|
||||
/// naturally; only the allocator's extent map is primed.
|
||||
pub(super) fn preallocate(file: &File, size_bytes: u64) {
|
||||
preallocate_impl(file, size_bytes);
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
//! Fallback preallocate impl. No-op.
|
||||
|
||||
use std::fs::File;
|
||||
|
||||
pub(super) fn preallocate_impl(_file: &File, size_bytes: u64) {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"LocalFileSink preallocate size_hint={size_bytes} skipped (no platform impl)"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
//! TCP / UDP socket sinks (sequential-only).
|
||||
//!
|
||||
//! [`SocketSink`] wraps a `TcpStream` in a 1 MiB `BufWriter`. Constructor
|
||||
//! tunes `SO_SNDBUF` to a caller hint when provided. `finish()` flushes
|
||||
//! the buffer then `shutdown(Write)`s the socket so the peer sees clean
|
||||
//! end-of-stream.
|
||||
//!
|
||||
//! [`UdpSocketSink`] wraps a connected `UdpSocket`. Each `write` call
|
||||
//! emits exactly one datagram — the caller is responsible for packetizing
|
||||
//! to a reasonable MTU (188 × 7 = 1316 bytes for MPEG-TS-over-UDP is the
|
||||
//! conventional choice). `finish()` is a no-op; UDP has no end-of-stream
|
||||
//! marker.
|
||||
//!
|
||||
//! Both types satisfy [`SequentialSink`] via the blanket impl in
|
||||
//! `super::mod`. Neither implements `Seek`, so neither satisfies
|
||||
//! [`RandomAccessSink`] — using one with `MkvMux` is a compile error,
|
||||
//! which is the design intent.
|
||||
//!
|
||||
//! [`SequentialSink`]: super::SequentialSink
|
||||
//! [`RandomAccessSink`]: super::RandomAccessSink
|
||||
|
||||
use std::io::{self, BufWriter, Write};
|
||||
use std::net::{Shutdown, TcpStream, ToSocketAddrs, UdpSocket};
|
||||
|
||||
/// `BufWriter` capacity for [`SocketSink`]. 1 MiB matches the typical
|
||||
/// kernel send-buffer ceiling and keeps small-write amplification from
|
||||
/// containers (TS = 188-byte packets, fMP4 fragment headers = ~100 bytes)
|
||||
/// from translating into syscall storms.
|
||||
const TCP_BUF_CAPACITY: usize = 1024 * 1024;
|
||||
|
||||
/// Sequential-only sink over a TCP connection.
|
||||
///
|
||||
/// Wraps a `BufWriter<TcpStream>`; the inner `TcpStream` is kept as a
|
||||
/// clone so [`finish`](Self::finish) can call `shutdown(Write)` after
|
||||
/// flushing the buffer (the buffered writer doesn't expose the socket
|
||||
/// directly).
|
||||
pub struct SocketSink {
|
||||
/// Buffered write half. All payload bytes go through this.
|
||||
buf: BufWriter<TcpStream>,
|
||||
/// Shutdown handle — clone of the socket inside `buf`. Used only by
|
||||
/// `finish()` for `shutdown(Write)`; never read or written through.
|
||||
shutdown_handle: TcpStream,
|
||||
}
|
||||
|
||||
impl SocketSink {
|
||||
/// Open a TCP connection to `addr` and wrap it for sequential
|
||||
/// writing. `sndbuf_bytes`, when present, is forwarded to
|
||||
/// `setsockopt(SO_SNDBUF)` as a kernel hint — the OS may clamp it.
|
||||
///
|
||||
/// `addr` accepts anything `ToSocketAddrs` does: `"10.0.0.1:1234"`,
|
||||
/// `("host", 1234)`, a `SocketAddr`, etc.
|
||||
pub fn connect<A: ToSocketAddrs>(addr: A, sndbuf_bytes: Option<usize>) -> io::Result<Self> {
|
||||
let stream = TcpStream::connect(addr)?;
|
||||
// `set_nodelay(true)` keeps small writes (TS packet trains, fMP4
|
||||
// moof headers) from sitting in Nagle's algorithm until the buffer
|
||||
// fills. The BufWriter already absorbs syscall overhead; Nagle
|
||||
// would just add latency without coalescing more.
|
||||
stream.set_nodelay(true)?;
|
||||
if let Some(n) = sndbuf_bytes {
|
||||
set_send_buffer(&stream, n)?;
|
||||
}
|
||||
let shutdown_handle = stream.try_clone()?;
|
||||
Ok(Self {
|
||||
buf: BufWriter::with_capacity(TCP_BUF_CAPACITY, stream),
|
||||
shutdown_handle,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for SocketSink {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.buf.write(buf)
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.buf.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl SocketSink {
|
||||
/// Drain the BufWriter and `shutdown(Write)` the underlying socket
|
||||
/// so the peer sees a clean EOF.
|
||||
///
|
||||
/// Note: [`SequentialSink::finish`](super::SequentialSink::finish)'s
|
||||
/// blanket-impl default is a no-op. Trait-object call sites that
|
||||
/// need socket shutdown should call this inherent method directly
|
||||
/// before dropping the sink, or hold the concrete `SocketSink` type
|
||||
/// (typical pattern: each muxer's `finish()` calls the appropriate
|
||||
/// inherent close method on its captured concrete sink).
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
self.buf.flush()?;
|
||||
// `shutdown(Write)` signals clean EOF to the peer. Errors here
|
||||
// are non-fatal — the connection may have already been torn down
|
||||
// by the peer — but we surface them so callers can log.
|
||||
self.shutdown_handle.shutdown(Shutdown::Write)
|
||||
}
|
||||
}
|
||||
|
||||
/// Sequential-only sink over a connected UDP socket.
|
||||
///
|
||||
/// Each [`write`](Write::write) call sends exactly one datagram. The
|
||||
/// caller is responsible for splitting payload at packet boundaries —
|
||||
/// for MPEG-TS this means 7 × 188 = 1316 bytes per datagram, the
|
||||
/// industry standard for MPEG-TS-over-UDP. No buffering happens here;
|
||||
/// adding it would silently merge datagrams.
|
||||
///
|
||||
/// `finish()` is a no-op: UDP has no end-of-stream marker. Closing the
|
||||
/// socket happens on drop.
|
||||
pub struct UdpSocketSink {
|
||||
socket: UdpSocket,
|
||||
}
|
||||
|
||||
impl UdpSocketSink {
|
||||
/// Bind a local UDP socket to an ephemeral port and `connect` it to
|
||||
/// `peer`. `connect` doesn't open a connection — it just fixes the
|
||||
/// peer address so subsequent `send` calls don't need to repeat it,
|
||||
/// and so receive-side filtering rejects packets from other sources.
|
||||
///
|
||||
/// `sndbuf_bytes`, when present, is a hint to `SO_SNDBUF`.
|
||||
pub fn connect<A: ToSocketAddrs>(peer: A, sndbuf_bytes: Option<usize>) -> io::Result<Self> {
|
||||
// Bind to all-zeros / any port. The kernel picks an ephemeral
|
||||
// source port and the source IP at first send.
|
||||
let socket = UdpSocket::bind("0.0.0.0:0")?;
|
||||
socket.connect(peer)?;
|
||||
if let Some(n) = sndbuf_bytes {
|
||||
set_udp_send_buffer(&socket, n)?;
|
||||
}
|
||||
Ok(Self { socket })
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for UdpSocketSink {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
// `send` writes the entire datagram or fails — no partial sends
|
||||
// for UDP. Match `Write::write`'s contract by reporting bytes
|
||||
// accepted.
|
||||
self.socket.send(buf)
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl UdpSocketSink {
|
||||
/// No-op — UDP has no end-of-stream marker. Provided for parity
|
||||
/// with [`SocketSink::finish`] so call sites can treat them uniformly.
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ── Platform `SO_SNDBUF` tuning ────────────────────────────────────────────
|
||||
//
|
||||
// std's `TcpStream` / `UdpSocket` don't expose `SO_SNDBUF`. We drop to
|
||||
// libc on Linux + macOS (the libc-dep targets in Cargo.toml). On other
|
||||
// targets the hint is silently ignored — the socket still works, the
|
||||
// kernel just picks its own send-buffer size.
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "macos"))]
|
||||
fn set_send_buffer(stream: &TcpStream, bytes: usize) -> io::Result<()> {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
setsockopt_sndbuf(stream.as_raw_fd(), bytes)
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "macos"))]
|
||||
fn set_udp_send_buffer(socket: &UdpSocket, bytes: usize) -> io::Result<()> {
|
||||
use std::os::unix::io::AsRawFd;
|
||||
setsockopt_sndbuf(socket.as_raw_fd(), bytes)
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
fn set_send_buffer(_stream: &TcpStream, _bytes: usize) -> io::Result<()> {
|
||||
// Non-Linux-non-macOS targets aren't in Cargo.toml's libc dep list;
|
||||
// silently ignore the hint rather than failing the connect. Callers
|
||||
// can detect via the lack of an explicit "sndbuf applied" signal
|
||||
// (not provided, intentionally — this is a hint, not a guarantee).
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos")))]
|
||||
fn set_udp_send_buffer(_socket: &UdpSocket, _bytes: usize) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(any(target_os = "linux", target_os = "macos"))]
|
||||
fn setsockopt_sndbuf(fd: std::os::unix::io::RawFd, bytes: usize) -> io::Result<()> {
|
||||
// Clamp into c_int range; SO_SNDBUF takes an `int` argument.
|
||||
let want: libc::c_int = bytes.try_into().unwrap_or(libc::c_int::MAX);
|
||||
let ret = unsafe {
|
||||
libc::setsockopt(
|
||||
fd,
|
||||
libc::SOL_SOCKET,
|
||||
libc::SO_SNDBUF,
|
||||
&want as *const _ as *const libc::c_void,
|
||||
std::mem::size_of::<libc::c_int>() as libc::socklen_t,
|
||||
)
|
||||
};
|
||||
if ret != 0 {
|
||||
return Err(io::Error::last_os_error());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Read;
|
||||
use std::net::{TcpListener, UdpSocket};
|
||||
use std::thread;
|
||||
|
||||
/// Bind a listener, accept on a thread, return (listener_addr,
|
||||
/// accepted-bytes future via JoinHandle).
|
||||
#[test]
|
||||
fn socket_sink_round_trips_bytes() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let accept = thread::spawn(move || {
|
||||
let (mut sock, _) = listener.accept().unwrap();
|
||||
let mut buf = Vec::new();
|
||||
sock.read_to_end(&mut buf).unwrap();
|
||||
buf
|
||||
});
|
||||
|
||||
let mut sink = SocketSink::connect(addr, Some(256 * 1024)).unwrap();
|
||||
// Write enough to overflow the BufWriter at least once, then a
|
||||
// tail that lives in the buffer until `finish` flushes.
|
||||
let big: Vec<u8> = (0..(2 * TCP_BUF_CAPACITY))
|
||||
.map(|i| (i & 0xff) as u8)
|
||||
.collect();
|
||||
sink.write_all(&big).unwrap();
|
||||
sink.write_all(b"tail\n").unwrap();
|
||||
sink.finish().unwrap();
|
||||
drop(sink);
|
||||
|
||||
let received = accept.join().unwrap();
|
||||
assert_eq!(received.len(), big.len() + 5);
|
||||
assert_eq!(&received[..big.len()], &big[..]);
|
||||
assert_eq!(&received[big.len()..], b"tail\n");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn socket_sink_is_sequential_only() {
|
||||
// Compile-time assertion via dyn — if this ever started
|
||||
// satisfying `RandomAccessSink`, the trait split would be broken.
|
||||
fn _assert_seq(_: &mut dyn super::super::SequentialSink) {}
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let _accept = thread::spawn(move || {
|
||||
let _ = listener.accept();
|
||||
});
|
||||
let mut sink = SocketSink::connect(addr, None).unwrap();
|
||||
_assert_seq(&mut sink);
|
||||
// The negative is harder to assert directly (no `is_not<T>`),
|
||||
// but `SocketSink` does not impl `Seek`, so it can't unify with
|
||||
// `RandomAccessSink`'s super-bound. The Phase 2 blanket impl
|
||||
// `impl<T: SequentialSink + Seek> RandomAccessSink for T {}` thus
|
||||
// excludes it by construction.
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn udp_socket_sink_delivers_datagrams() {
|
||||
let receiver = UdpSocket::bind("127.0.0.1:0").unwrap();
|
||||
receiver
|
||||
.set_read_timeout(Some(std::time::Duration::from_secs(2)))
|
||||
.unwrap();
|
||||
let addr = receiver.local_addr().unwrap();
|
||||
let mut sink = UdpSocketSink::connect(addr, Some(128 * 1024)).unwrap();
|
||||
|
||||
sink.write_all(&[1, 2, 3, 4, 5]).unwrap();
|
||||
sink.write_all(&[9, 9, 9]).unwrap();
|
||||
sink.finish().unwrap();
|
||||
|
||||
let mut buf = [0u8; 64];
|
||||
let n1 = receiver.recv(&mut buf).unwrap();
|
||||
assert_eq!(&buf[..n1], &[1, 2, 3, 4, 5]);
|
||||
let n2 = receiver.recv(&mut buf).unwrap();
|
||||
assert_eq!(&buf[..n2], &[9, 9, 9]);
|
||||
}
|
||||
}
|
||||
+129
-24
@@ -300,31 +300,22 @@ impl WritebackPipeline {
|
||||
}
|
||||
}
|
||||
|
||||
/// Probe whether `fd` lives on an NFS mount via `fstatfs`. Returns
|
||||
/// `false` on any error — we fail open, not closed: better to run the
|
||||
/// normal local-storage path on a misdetected NFS mount (and surface
|
||||
/// the freeze loudly via the timeout) than to needlessly disable
|
||||
/// writeback bounding on every local file because of a transient
|
||||
/// stat error.
|
||||
/// Probe whether `fd` lives on an NFS mount. Thin wrapper around
|
||||
/// [`crate::platform::fs_type::detect_fd`] so writeback policy and
|
||||
/// general-purpose fs-type classification stay in sync (same magic
|
||||
/// numbers, same musl-vs-glibc cast handling).
|
||||
///
|
||||
/// Fails open: any classification other than NFS counts as "not NFS"
|
||||
/// (including `Unknown` on `fstatfs` error) — better to run the
|
||||
/// normal local-storage path on a misdetected NFS mount and surface
|
||||
/// the freeze loudly via [`WAIT_AFTER_TIMEOUT`] than to needlessly
|
||||
/// disable writeback bounding on every local file because of a
|
||||
/// transient stat error.
|
||||
fn detect_nfs(fd: RawFd) -> bool {
|
||||
// `libc::statfs` is repr(C) with a fixed layout; zeroing is the
|
||||
// documented init pattern for the kernel uapi struct.
|
||||
let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::fstatfs(fd, &mut buf) };
|
||||
if rc != 0 {
|
||||
let errno = std::io::Error::last_os_error();
|
||||
tracing::warn!(
|
||||
target: "mux",
|
||||
"WritebackPipeline fstatfs(fd={fd}) failed: {errno} — defaulting is_nfs=false",
|
||||
);
|
||||
return false;
|
||||
}
|
||||
// `f_type` is signed (`__fsword_t`) on glibc and unsigned
|
||||
// (`c_ulong`) on musl. Cast both sides to i64 for a portable
|
||||
// comparison.
|
||||
let f_type = buf.f_type as i64;
|
||||
let nfs_magic = libc::NFS_SUPER_MAGIC as i64;
|
||||
f_type == nfs_magic
|
||||
matches!(
|
||||
crate::platform::fs_type::detect_fd(fd),
|
||||
crate::platform::fs_type::FsType::Nfs
|
||||
)
|
||||
}
|
||||
|
||||
/// Run `sync_file_range(WAIT_AFTER)` on a worker thread and wait up
|
||||
@@ -352,3 +343,117 @@ fn wait_after_with_timeout(fd: RawFd, off: u64, len: u64) -> Option<u64> {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
/// Helper: build a `WritebackPipeline` over a local tempfile. On
|
||||
/// every test rig (linux dev box, CI) the tempfile lives on a
|
||||
/// local FS, so `is_nfs=false` and `skip_wait` returns false until
|
||||
/// we explicitly mark the pipeline degraded.
|
||||
fn local_pipeline(chunk_bytes: u64) -> (NamedTempFile, WritebackPipeline) {
|
||||
let f = NamedTempFile::new().expect("tempfile create");
|
||||
let pipeline = WritebackPipeline::new(f.as_file(), 0, chunk_bytes);
|
||||
(f, pipeline)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn new_pipeline_starts_active() {
|
||||
let (_f, p) = local_pipeline(32 * 1024 * 1024);
|
||||
assert!(!p.is_nfs, "local tempfile must not classify as NFS");
|
||||
assert!(!p.degraded.load(Ordering::Relaxed));
|
||||
assert!(!p.skip_wait(), "fresh local pipeline must not skip wait");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn degraded_flag_short_circuits_wait() {
|
||||
let (_f, p) = local_pipeline(32 * 1024 * 1024);
|
||||
assert!(!p.skip_wait());
|
||||
p.degraded.store(true, Ordering::Relaxed);
|
||||
assert!(
|
||||
p.skip_wait(),
|
||||
"degraded flag must force the wait+dontneed bypass"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_grows_chunk_on_high_p95() {
|
||||
let (_f, mut p) = local_pipeline(16 * 1024 * 1024);
|
||||
// Fill the window with samples above the grow threshold.
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(ADAPTIVE_GROW_MS + 50);
|
||||
}
|
||||
assert!(
|
||||
p.chunk_bytes > 16 * 1024 * 1024,
|
||||
"chunk should have grown; got {}",
|
||||
p.chunk_bytes
|
||||
);
|
||||
assert!(p.chunk_bytes <= CHUNK_BYTES_MAX);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_shrinks_chunk_on_low_p95() {
|
||||
let (_f, mut p) = local_pipeline(64 * 1024 * 1024);
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(1); // well under ADAPTIVE_SHRINK_MS
|
||||
}
|
||||
assert!(
|
||||
p.chunk_bytes < 64 * 1024 * 1024,
|
||||
"chunk should have shrunk; got {}",
|
||||
p.chunk_bytes
|
||||
);
|
||||
assert!(p.chunk_bytes >= CHUNK_BYTES_MIN);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_no_op_below_window_fill() {
|
||||
let (_f, mut p) = local_pipeline(16 * 1024 * 1024);
|
||||
let initial = p.chunk_bytes;
|
||||
// Only push a few samples; window not full → no adaptation.
|
||||
for _ in 0..(ADAPTIVE_WINDOW - 1) {
|
||||
p.record_wait(ADAPTIVE_GROW_MS + 100);
|
||||
}
|
||||
assert_eq!(
|
||||
p.chunk_bytes, initial,
|
||||
"chunk must not change before window is full"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn record_wait_clamps_to_chunk_bounds() {
|
||||
// Grow past the max.
|
||||
let (_f, mut p) = local_pipeline(CHUNK_BYTES_MAX);
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(ADAPTIVE_GROW_MS + 1000);
|
||||
}
|
||||
assert_eq!(p.chunk_bytes, CHUNK_BYTES_MAX, "must clamp to MAX");
|
||||
|
||||
// Shrink past the min.
|
||||
let (_f, mut p) = local_pipeline(CHUNK_BYTES_MIN);
|
||||
for _ in 0..ADAPTIVE_WINDOW {
|
||||
p.record_wait(0);
|
||||
}
|
||||
assert_eq!(p.chunk_bytes, CHUNK_BYTES_MIN, "must clamp to MIN");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detect_nfs_local_file_is_false() {
|
||||
// Local tempfile must not classify as NFS. This locks in the
|
||||
// consolidation through `crate::platform::fs_type::detect_fd`.
|
||||
let f = NamedTempFile::new().expect("tempfile create");
|
||||
use std::os::unix::io::AsRawFd;
|
||||
assert!(!detect_nfs(f.as_file().as_raw_fd()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn note_progress_below_chunk_is_noop() {
|
||||
let (_f, mut p) = local_pipeline(32 * 1024 * 1024);
|
||||
// No-op return before crossing the first chunk boundary.
|
||||
let before = p.chunk_count;
|
||||
p.note_progress(1024); // < 32 MiB
|
||||
assert_eq!(p.chunk_count, before);
|
||||
assert!(p.pending.is_none());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,229 +0,0 @@
|
||||
//! `WritebackFile` — a `File` wrapper whose reason for existing is the
|
||||
//! bounded-cache writeback pipeline.
|
||||
//!
|
||||
//! Why: large sequential writes (sweep, patch, mux on UHD-scale output)
|
||||
//! left to the kernel's default writeback policy accumulate hundreds of
|
||||
//! megabytes of dirty pages and then burst-flush, stalling subsequent
|
||||
//! writes for seconds at a time. `WritebackFile` drives a continuous
|
||||
//! [`super::writeback::WritebackPipeline`] that on Linux issues
|
||||
//! incremental `sync_file_range` + `posix_fadvise(DONTNEED)` calls at
|
||||
//! 32 MB granularity so dirty pages drain at the same rate they're
|
||||
//! produced. macOS and Windows fall through to a no-op pipeline — their
|
||||
//! default cache policies have not been shown to exhibit the same
|
||||
//! pathology for this access pattern.
|
||||
//!
|
||||
//! It implements `Write` and `Seek` so any call site that wrote to a
|
||||
//! plain `File` through those traits (sweep, patch, mux) can swap in
|
||||
//! `WritebackFile` without touching the body of the loop. The wrapper
|
||||
//! also tracks the current file position to feed the pipeline with
|
||||
//! progress + seek boundaries.
|
||||
//!
|
||||
//! See `super::writeback::linux` for the underlying pathology and the
|
||||
//! strategy.
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{self, Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use super::writeback::WritebackPipeline;
|
||||
|
||||
const WRITEBACK_CHUNK_BYTES: u64 = 32 * 1024 * 1024;
|
||||
|
||||
pub(crate) struct WritebackFile {
|
||||
file: File,
|
||||
pipeline: WritebackPipeline,
|
||||
pos: u64,
|
||||
}
|
||||
|
||||
impl WritebackFile {
|
||||
/// Wrap an open `File`. The current OS file position is queried
|
||||
/// once so the pipeline starts tracking from wherever the file
|
||||
/// already is (typically 0 for fresh files; non-zero for resumed
|
||||
/// or appended files).
|
||||
pub(crate) fn new(mut file: File) -> io::Result<Self> {
|
||||
let pos = file.stream_position()?;
|
||||
let pipeline = WritebackPipeline::new(&file, pos, WRITEBACK_CHUNK_BYTES);
|
||||
Ok(Self {
|
||||
file,
|
||||
pipeline,
|
||||
pos,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new file at `path` (truncating any existing contents)
|
||||
/// and wrap it. Convenience for the common
|
||||
/// `File::create(path)` + `WritebackFile::new(file)` pair so callers
|
||||
/// don't have to assemble a `File` first.
|
||||
///
|
||||
/// Callers that know the target output size should prefer
|
||||
/// [`Self::create_with_size_hint`] so the kernel can pre-reserve
|
||||
/// extents.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn create(path: &Path) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Like [`Self::create`] but pre-reserves `size_bytes` of disk
|
||||
/// space via `fallocate(FALLOC_FL_KEEP_SIZE)` on Linux. The
|
||||
/// reported file size is unchanged (writes still grow the file
|
||||
/// naturally) — only the on-disk extent allocation is preallocated,
|
||||
/// which reduces extent fragmentation on large sequential writes
|
||||
/// (mux output, especially on slow storage / NFS).
|
||||
///
|
||||
/// On macOS / Windows the size hint is ignored and this is
|
||||
/// equivalent to `create`.
|
||||
pub(crate) fn create_with_size_hint(path: &Path, size_bytes: u64) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
use std::os::unix::io::AsRawFd;
|
||||
// FALLOC_FL_KEEP_SIZE = 0x01 — keep the reported file
|
||||
// size at 0 (writes grow it normally) while still
|
||||
// pre-reserving the extents.
|
||||
let rc = unsafe {
|
||||
libc::fallocate(
|
||||
file.as_raw_fd(),
|
||||
libc::FALLOC_FL_KEEP_SIZE,
|
||||
0,
|
||||
size_bytes as i64,
|
||||
)
|
||||
};
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile fallocate size_hint={size_bytes} rc={rc} ok={}",
|
||||
rc == 0
|
||||
);
|
||||
}
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
{
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile fallocate size_hint={size_bytes} skipped (non-linux)"
|
||||
);
|
||||
}
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Open an existing file at `path` for writing (no truncation) and
|
||||
/// wrap it. Mirrors `File::open` semantics for the writable case
|
||||
/// — used by patch / resume paths that mutate an existing ISO in
|
||||
/// place.
|
||||
pub(crate) fn open(path: &Path) -> io::Result<Self> {
|
||||
let file = OpenOptions::new().write(true).open(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Drain in-flight writeback then issue a full fsync. Use this in
|
||||
/// place of `File::sync_all`.
|
||||
///
|
||||
/// The final fsync is wrapped in
|
||||
/// [`crate::io::bounded::bounded_syscall`] with a 60 s deadline.
|
||||
/// fsync on a wedged NFS server (or a degraded local disk) can
|
||||
/// hang the calling thread; the wrapper ensures the worst case is
|
||||
/// 60 s + log-and-continue rather than indefinite. On timeout the
|
||||
/// page cache is left to the kernel's normal flush-on-close path —
|
||||
/// best effort, but bounded.
|
||||
pub(crate) fn sync_all(&mut self) -> io::Result<()> {
|
||||
self.pipeline.finalize();
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use std::time::Duration;
|
||||
let fd = self.file.as_raw_fd();
|
||||
match crate::io::bounded::bounded_syscall(
|
||||
None,
|
||||
Duration::from_secs(60),
|
||||
move || -> io::Result<()> {
|
||||
let rc = unsafe { libc::fsync(fd) };
|
||||
if rc == 0 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::last_os_error())
|
||||
}
|
||||
},
|
||||
) {
|
||||
Ok(inner) => inner,
|
||||
Err(crate::io::bounded::BoundedError::Timeout) => {
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackFile::sync_all fsync timed out after 60s; kernel will flush on close (best-effort)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(crate::io::bounded::BoundedError::Halted) => Ok(()),
|
||||
Err(crate::io::bounded::BoundedError::WorkerLost) => Ok(()),
|
||||
}
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
{
|
||||
self.file.sync_all()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for WritebackFile {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
let n = self.file.write(buf)?;
|
||||
self.pos += n as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.file.write_all(buf)?;
|
||||
self.pos += buf.len() as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.file.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl Seek for WritebackFile {
|
||||
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
|
||||
let p = self.file.seek(from)?;
|
||||
// Only treat seeks that actually move the position as
|
||||
// boundaries — sweep does a redundant `seek(Current(pos))`
|
||||
// before every write, and we don't want that to drain the
|
||||
// pipeline on every iteration.
|
||||
if p != self.pos {
|
||||
// Diagnostic for the NFS 73 % mux hang: the MKV format
|
||||
// requires the muxer to seek back occasionally (cluster
|
||||
// size patching, Cues index write, Segment header
|
||||
// backpatch). Each such seek invalidates the writeback
|
||||
// chunk tracking and forces a finalize → WAIT_AFTER on
|
||||
// the in-flight chunk. Logging the seek delta lets us
|
||||
// correlate hang offsets with specific muxer operations.
|
||||
let from_pos = self.pos;
|
||||
let to_pos = p;
|
||||
let delta: i64 = (to_pos as i64).wrapping_sub(from_pos as i64);
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile seek from={from_pos} to={to_pos} delta={delta}"
|
||||
);
|
||||
self.pipeline.handle_seek(p);
|
||||
self.pos = p;
|
||||
}
|
||||
Ok(p)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WritebackFile {
|
||||
fn drop(&mut self) {
|
||||
// Run the pipeline's tail finalize so the last in-flight chunk
|
||||
// gets its `WAIT_AFTER` + `posix_fadvise(DONTNEED)`. Without
|
||||
// this, callers that drop a `WritebackFile` without calling
|
||||
// `sync_all` (panic, early-return, idiomatic `let _ = w;`)
|
||||
// leave the trailing chunk in cache; the kernel still flushes
|
||||
// on close, but the bounded-cache invariant fails at the tail.
|
||||
// We deliberately do *not* call `self.file.sync_all()` here —
|
||||
// close already triggers a flush, and an `fsync` from `Drop`
|
||||
// would silently swallow its `io::Error` anyway. `finalize` is
|
||||
// idempotent so an explicit `sync_all` followed by drop is
|
||||
// still safe.
|
||||
self.pipeline.finalize();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
//! Linux platform impl for [`super::WritebackFile`].
|
||||
//!
|
||||
//! - `preallocate`: `fallocate(FALLOC_FL_KEEP_SIZE)` — reserve extents
|
||||
//! without growing the reported file size. Reduces extent
|
||||
//! fragmentation on large sequential writes (mux output on NFS in
|
||||
//! particular).
|
||||
//! - `durable_sync`: `fsync` wrapped in
|
||||
//! [`crate::io::bounded::bounded_syscall`] with a 60 s deadline so a
|
||||
//! wedged NFS server can't trap the calling thread indefinitely.
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Pre-reserve extents for `size_bytes` of upcoming sequential writes.
|
||||
/// Best-effort: a non-zero rc is logged but not propagated, since the
|
||||
/// caller would just continue with the unreserved file anyway.
|
||||
pub(super) fn preallocate(file: &File, size_bytes: u64) {
|
||||
// FALLOC_FL_KEEP_SIZE = 0x01 — keep the reported file size at 0
|
||||
// (writes grow it normally) while still pre-reserving the extents.
|
||||
let rc = unsafe {
|
||||
libc::fallocate(
|
||||
file.as_raw_fd(),
|
||||
libc::FALLOC_FL_KEEP_SIZE,
|
||||
0,
|
||||
size_bytes as i64,
|
||||
)
|
||||
};
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile fallocate size_hint={size_bytes} rc={rc} ok={}",
|
||||
rc == 0
|
||||
);
|
||||
}
|
||||
|
||||
/// Run `fsync` on `file` with a 60 s deadline. On timeout we log loudly
|
||||
/// and return `Ok(())` — the kernel will still flush on close, so the
|
||||
/// data is best-effort durable; the alternative (trap the thread for
|
||||
/// the rest of the rip) defeats `/api/stop`.
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
let fd = file.as_raw_fd();
|
||||
match crate::io::bounded::bounded_syscall(
|
||||
None,
|
||||
Duration::from_secs(60),
|
||||
move || -> io::Result<()> {
|
||||
let rc = unsafe { libc::fsync(fd) };
|
||||
if rc == 0 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::last_os_error())
|
||||
}
|
||||
},
|
||||
) {
|
||||
Ok(inner) => inner,
|
||||
Err(crate::io::bounded::BoundedError::Timeout) => {
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackFile::sync_all fsync timed out after 60s; kernel will flush on close (best-effort)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(crate::io::bounded::BoundedError::Halted) => Ok(()),
|
||||
Err(crate::io::bounded::BoundedError::WorkerLost) => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
//! macOS platform impl for [`super::WritebackFile`].
|
||||
//!
|
||||
//! - `preallocate`: `fcntl(F_PREALLOCATE)` — macOS's fallocate-equiv.
|
||||
//! Reserves a contiguous extent when possible, falling back to a
|
||||
//! non-contiguous reservation if the FS can't satisfy it. Reported
|
||||
//! file size is unchanged (`F_ALLOCATEALL` is not set, so allocation
|
||||
//! is "best effort up to length"; growth happens via writes).
|
||||
//! - `durable_sync`: `fcntl(F_FULLFSYNC)` wrapped in
|
||||
//! [`crate::io::bounded::bounded_syscall`] with a 60 s deadline.
|
||||
//! F_FULLFSYNC is HFS+/APFS's true-fsync (flushes the disk's own
|
||||
//! write cache) — what `fsync` should have been on macOS. Falls back
|
||||
//! to plain `fsync` if F_FULLFSYNC returns ENOTSUP.
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
use std::os::unix::io::AsRawFd;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::io::platform_macos::{
|
||||
F_ALLOCATEALL, F_ALLOCATECONTIG, F_PEOFPOSMODE, F_PREALLOCATE, Fstore,
|
||||
};
|
||||
|
||||
/// `fcntl(F_FULLFSYNC)` opcode. Documented in `man 2 fcntl` on macOS;
|
||||
/// not in the `libc` crate as a named constant.
|
||||
const F_FULLFSYNC: libc::c_int = 51;
|
||||
|
||||
pub(super) fn preallocate(file: &File, size_bytes: u64) {
|
||||
let mut fst = Fstore {
|
||||
fst_flags: F_ALLOCATECONTIG | F_ALLOCATEALL,
|
||||
fst_posmode: F_PEOFPOSMODE,
|
||||
fst_offset: 0,
|
||||
fst_length: size_bytes as libc::off_t,
|
||||
fst_bytesalloc: 0,
|
||||
};
|
||||
// First attempt: contiguous.
|
||||
let mut rc = unsafe { libc::fcntl(file.as_raw_fd(), F_PREALLOCATE, &mut fst) };
|
||||
if rc == -1 {
|
||||
// Fall back: drop the contiguous hint, allow scattered extents.
|
||||
fst.fst_flags = F_ALLOCATEALL;
|
||||
rc = unsafe { libc::fcntl(file.as_raw_fd(), F_PREALLOCATE, &mut fst) };
|
||||
}
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile F_PREALLOCATE size_hint={size_bytes} rc={rc} bytes_allocated={} ok={}",
|
||||
fst.fst_bytesalloc,
|
||||
rc != -1
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
let fd = file.as_raw_fd();
|
||||
match crate::io::bounded::bounded_syscall(
|
||||
None,
|
||||
Duration::from_secs(60),
|
||||
move || -> io::Result<()> {
|
||||
// Try F_FULLFSYNC first. If it isn't supported on this
|
||||
// filesystem (older HFS, some network mounts) fall back to
|
||||
// plain fsync — better than nothing.
|
||||
let rc = unsafe { libc::fcntl(fd, F_FULLFSYNC, 0) };
|
||||
if rc == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let err = io::Error::last_os_error();
|
||||
if err.raw_os_error() == Some(libc::ENOTSUP) {
|
||||
let rc = unsafe { libc::fsync(fd) };
|
||||
if rc == 0 {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(io::Error::last_os_error())
|
||||
}
|
||||
} else {
|
||||
Err(err)
|
||||
}
|
||||
},
|
||||
) {
|
||||
Ok(inner) => inner,
|
||||
Err(crate::io::bounded::BoundedError::Timeout) => {
|
||||
tracing::error!(
|
||||
target: "mux",
|
||||
"WritebackFile::sync_all F_FULLFSYNC timed out after 60s; kernel will flush on close (best-effort)"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
Err(crate::io::bounded::BoundedError::Halted) => Ok(()),
|
||||
Err(crate::io::bounded::BoundedError::WorkerLost) => Ok(()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,315 @@
|
||||
//! `WritebackFile` — a `File` wrapper whose reason for existing is the
|
||||
//! bounded-cache writeback pipeline.
|
||||
//!
|
||||
//! Why: large sequential writes (sweep, patch, mux on UHD-scale output)
|
||||
//! left to the kernel's default writeback policy accumulate hundreds of
|
||||
//! megabytes of dirty pages and then burst-flush, stalling subsequent
|
||||
//! writes for seconds at a time. `WritebackFile` drives a continuous
|
||||
//! [`super::writeback::WritebackPipeline`] that on Linux issues
|
||||
//! incremental `sync_file_range` + `posix_fadvise(DONTNEED)` calls at
|
||||
//! 32 MB granularity so dirty pages drain at the same rate they're
|
||||
//! produced. macOS and Windows fall through to a no-op pipeline — their
|
||||
//! default cache policies have not been shown to exhibit the same
|
||||
//! pathology for this access pattern.
|
||||
//!
|
||||
//! It implements `Write` and `Seek` so any call site that wrote to a
|
||||
//! plain `File` through those traits (sweep, patch, mux) can swap in
|
||||
//! `WritebackFile` without touching the body of the loop. The wrapper
|
||||
//! also tracks the current file position to feed the pipeline with
|
||||
//! progress + seek boundaries.
|
||||
//!
|
||||
//! See `super::writeback::linux` for the underlying pathology and the
|
||||
//! strategy.
|
||||
//!
|
||||
//! ## Platform split
|
||||
//!
|
||||
//! The platform-specific pieces of this wrapper — extent preallocation
|
||||
//! (Linux `fallocate(KEEP_SIZE)`, macOS `F_PREALLOCATE`, Windows
|
||||
//! `SetFileValidData`) and the durable-flush primitive (Linux/macOS
|
||||
//! `fsync`/`F_FULLFSYNC` wrapped in a bounded syscall, Windows
|
||||
//! `FlushFileBuffers`) — live in per-OS sibling modules. The dispatch
|
||||
//! happens once at the bottom of this file via cfg-gated `mod` decls.
|
||||
//! No inline `#[cfg(target_os = "...")]` in the business-logic above.
|
||||
//!
|
||||
//! ## Write path
|
||||
//!
|
||||
//! Writes are direct passthrough to the underlying `File` (no writer
|
||||
//! thread, no ring, no batching). Empirically the Phase-2.5
|
||||
//! writer-thread architecture introduced a ~60% mux throughput
|
||||
//! regression on NFS bidirectional workloads; reverting the write path
|
||||
//! to direct passthrough restores the 0.20.7 baseline. The writeback
|
||||
//! pipeline still runs (it's called inline from `write` / `write_all` /
|
||||
//! `seek`) so the bounded-cache invariant on Linux is preserved.
|
||||
//!
|
||||
//! ## Halt-safety
|
||||
//!
|
||||
//! `sync_all` runs the per-OS durable-flush primitive, which on
|
||||
//! Linux/macOS is wrapped in [`crate::io::bounded::bounded_syscall`]
|
||||
//! with a 60 s deadline. A wedged NFS server cannot trap the muxer
|
||||
//! indefinitely on the final fsync.
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
mod other;
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux as platform;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos as platform;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
use other as platform;
|
||||
#[cfg(target_os = "windows")]
|
||||
use windows as platform;
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{self, Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use super::writeback::WritebackPipeline;
|
||||
|
||||
/// Granularity at which the Linux writeback pipeline issues
|
||||
/// `sync_file_range` pairs. 32 MiB is the empirically best value on
|
||||
/// the rip1 test bed (NFS to unraid-1 over 1 GbE, single-disk SAS):
|
||||
/// 8 MiB / 64 MiB / 128 MiB all measured worse in the 0.21.x mux
|
||||
/// iteration runs. Override via `FREEMKV_WRITEBACK_CHUNK_MIB` —
|
||||
/// faster backends (NVMe, RAID) may tolerate larger windows.
|
||||
const WRITEBACK_CHUNK_BYTES_DEFAULT: u64 = 32 * 1024 * 1024;
|
||||
|
||||
fn writeback_chunk_bytes() -> u64 {
|
||||
std::env::var("FREEMKV_WRITEBACK_CHUNK_MIB")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u64>().ok())
|
||||
.filter(|&n| n > 0)
|
||||
.map(|n| n * 1024 * 1024)
|
||||
.unwrap_or(WRITEBACK_CHUNK_BYTES_DEFAULT)
|
||||
}
|
||||
|
||||
pub(crate) struct WritebackFile {
|
||||
file: File,
|
||||
pipeline: WritebackPipeline,
|
||||
pos: u64,
|
||||
}
|
||||
|
||||
impl WritebackFile {
|
||||
/// Wrap an open `File`. The current OS file position is queried
|
||||
/// once so the pipeline starts tracking from wherever the file
|
||||
/// already is (typically 0 for fresh files; non-zero for resumed
|
||||
/// or appended files).
|
||||
pub(crate) fn new(mut file: File) -> io::Result<Self> {
|
||||
let pos = file.stream_position()?;
|
||||
let pipeline = WritebackPipeline::new(&file, pos, writeback_chunk_bytes());
|
||||
Ok(Self {
|
||||
file,
|
||||
pipeline,
|
||||
pos,
|
||||
})
|
||||
}
|
||||
|
||||
/// Create a new file at `path` (truncating any existing contents)
|
||||
/// and wrap it. Convenience for the common
|
||||
/// `File::create(path)` + `WritebackFile::new(file)` pair so callers
|
||||
/// don't have to assemble a `File` first.
|
||||
///
|
||||
/// Callers that know the target output size should prefer
|
||||
/// [`Self::create_with_size_hint`] so the kernel can pre-reserve
|
||||
/// extents.
|
||||
#[allow(dead_code)]
|
||||
pub(crate) fn create(path: &Path) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Like [`Self::create`] but pre-reserves `size_bytes` of disk
|
||||
/// space via the platform's extent-preallocation primitive (Linux
|
||||
/// `fallocate(KEEP_SIZE)`, macOS `F_PREALLOCATE`, Windows
|
||||
/// `SetFileValidData` stub). The reported file size is unchanged
|
||||
/// (writes still grow the file naturally) — only the on-disk extent
|
||||
/// allocation is preallocated, which reduces extent fragmentation
|
||||
/// on large sequential writes (mux output, especially on slow
|
||||
/// storage / NFS).
|
||||
///
|
||||
/// On platforms without an extent-preallocation primitive this is
|
||||
/// equivalent to `create` — the size hint is dropped after a debug
|
||||
/// log.
|
||||
pub(crate) fn create_with_size_hint(path: &Path, size_bytes: u64) -> io::Result<Self> {
|
||||
let file = File::create(path)?;
|
||||
platform::preallocate(&file, size_bytes);
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Open an existing file at `path` for writing (no truncation) and
|
||||
/// wrap it. Mirrors `File::open` semantics for the writable case
|
||||
/// — used by patch / resume paths that mutate an existing ISO in
|
||||
/// place.
|
||||
pub(crate) fn open(path: &Path) -> io::Result<Self> {
|
||||
let file = OpenOptions::new().write(true).open(path)?;
|
||||
Self::new(file)
|
||||
}
|
||||
|
||||
/// Drain in-flight writeback then issue a full fsync. Use this in
|
||||
/// place of `File::sync_all`.
|
||||
///
|
||||
/// The final durable flush is wrapped in
|
||||
/// [`crate::io::bounded::bounded_syscall`] (per the per-OS module)
|
||||
/// with a 60 s deadline on Linux/macOS — a wedged NFS server cannot
|
||||
/// trap the calling thread indefinitely. On timeout the page cache
|
||||
/// is left to the kernel's normal flush-on-close path — best
|
||||
/// effort, but bounded.
|
||||
pub(crate) fn sync_all(&mut self) -> io::Result<()> {
|
||||
self.pipeline.finalize();
|
||||
platform::durable_sync(&self.file)
|
||||
}
|
||||
}
|
||||
|
||||
impl Write for WritebackFile {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
let n = self.file.write(buf)?;
|
||||
self.pos += n as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(n)
|
||||
}
|
||||
|
||||
fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
|
||||
self.file.write_all(buf)?;
|
||||
self.pos += buf.len() as u64;
|
||||
self.pipeline.note_progress(self.pos);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.file.flush()
|
||||
}
|
||||
}
|
||||
|
||||
impl Seek for WritebackFile {
|
||||
fn seek(&mut self, from: SeekFrom) -> io::Result<u64> {
|
||||
let p = self.file.seek(from)?;
|
||||
// Only treat seeks that actually move the position as
|
||||
// boundaries — sweep does a redundant `seek(Current(pos))`
|
||||
// before every write, and we don't want that to drain the
|
||||
// pipeline on every iteration.
|
||||
if p != self.pos {
|
||||
// Diagnostic for the NFS mux hang: the MKV format requires
|
||||
// the muxer to seek back occasionally (cluster size
|
||||
// patching, Cues index write, Segment header backpatch).
|
||||
// Each such seek invalidates the writeback chunk tracking
|
||||
// and forces a finalize → WAIT_AFTER on the in-flight
|
||||
// chunk. Logging the seek delta lets us correlate hang
|
||||
// offsets with specific muxer operations.
|
||||
let from_pos = self.pos;
|
||||
let to_pos = p;
|
||||
let delta: i64 = (to_pos as i64).wrapping_sub(from_pos as i64);
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile seek from={from_pos} to={to_pos} delta={delta}"
|
||||
);
|
||||
self.pipeline.handle_seek(p);
|
||||
self.pos = p;
|
||||
}
|
||||
Ok(p)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for WritebackFile {
|
||||
fn drop(&mut self) {
|
||||
// Run the pipeline's tail finalize so the last in-flight chunk
|
||||
// gets its `WAIT_AFTER` + `posix_fadvise(DONTNEED)`. Without
|
||||
// this, callers that drop a `WritebackFile` without calling
|
||||
// `sync_all` (panic, early-return, idiomatic `let _ = w;`)
|
||||
// leave the trailing chunk in cache; the kernel still flushes
|
||||
// on close, but the bounded-cache invariant fails at the tail.
|
||||
// We deliberately do *not* call `self.file.sync_all()` here —
|
||||
// close already triggers a flush, and an `fsync` from `Drop`
|
||||
// would silently swallow its `io::Error` anyway. `finalize` is
|
||||
// idempotent so an explicit `sync_all` followed by drop is
|
||||
// still safe.
|
||||
self.pipeline.finalize();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Read;
|
||||
|
||||
fn read_back(path: &Path) -> Vec<u8> {
|
||||
let mut f = File::open(path).unwrap();
|
||||
let mut v = Vec::new();
|
||||
f.read_to_end(&mut v).unwrap();
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_then_drop_persists_bytes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("a.bin");
|
||||
{
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
w.write_all(b"hello world").unwrap();
|
||||
// Drop drains the pipeline tail.
|
||||
}
|
||||
assert_eq!(read_back(&p), b"hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_all_drains_and_flushes() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("b.bin");
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
for _ in 0..32 {
|
||||
w.write_all(&[0x5au8; 1024]).unwrap();
|
||||
}
|
||||
// After sync_all, the bytes MUST be visible to a separate
|
||||
// reader. The pipeline has been finalised and durable-sync has
|
||||
// run.
|
||||
w.sync_all().unwrap();
|
||||
let bytes = read_back(&p);
|
||||
assert_eq!(bytes.len(), 32 * 1024);
|
||||
assert!(bytes.iter().all(|&b| b == 0x5a));
|
||||
drop(w);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seek_then_patch_roundtrip() {
|
||||
// Write A; seek back; patch with B; read back; the patch lands
|
||||
// at the right offset.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("c.bin");
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
let big = vec![b'A'; 4096];
|
||||
w.write_all(&big).unwrap();
|
||||
// Seek back to offset 1000 and overwrite 8 bytes.
|
||||
w.seek(SeekFrom::Start(1000)).unwrap();
|
||||
w.write_all(b"PATCHED!").unwrap();
|
||||
w.sync_all().unwrap();
|
||||
drop(w);
|
||||
let bytes = read_back(&p);
|
||||
assert_eq!(bytes.len(), 4096);
|
||||
assert_eq!(&bytes[1000..1008], b"PATCHED!");
|
||||
// Bytes outside the patch are still 'A'.
|
||||
assert_eq!(bytes[999], b'A');
|
||||
assert_eq!(bytes[1008], b'A');
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flush_is_observed_in_order() {
|
||||
// `Write::flush` should not panic or reorder; verify the bytes
|
||||
// land in order through interleaved flushes.
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let p = dir.path().join("f.bin");
|
||||
let mut w = WritebackFile::create(&p).unwrap();
|
||||
w.write_all(b"one").unwrap();
|
||||
w.flush().unwrap();
|
||||
w.write_all(b"two").unwrap();
|
||||
w.flush().unwrap();
|
||||
w.write_all(b"three").unwrap();
|
||||
w.sync_all().unwrap();
|
||||
drop(w);
|
||||
assert_eq!(read_back(&p), b"onetwothree");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
//! Fallback platform impl for [`super::WritebackFile`] on targets
|
||||
//! without a dedicated implementation (BSDs, illumos, etc.).
|
||||
//!
|
||||
//! - `preallocate` is a logged no-op.
|
||||
//! - `durable_sync` calls `File::sync_all` directly (no bounded-syscall
|
||||
//! wrapper — the wrapper depends on Linux/macOS unix idioms that
|
||||
//! aren't universally portable). If a future BSD impl needs the
|
||||
//! 60-s deadline, it should land in its own per-OS file rather than
|
||||
//! bloat this fallback.
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
|
||||
pub(super) fn preallocate(_file: &File, size_bytes: u64) {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile preallocate size_hint={size_bytes} skipped (no impl on this target)"
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
file.sync_all()
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
//! Windows platform impl for [`super::WritebackFile`].
|
||||
//!
|
||||
//! TODO: this stub matches the design's "validate without a Windows
|
||||
//! build env, leave a stub" carve-out. The real impl should use:
|
||||
//!
|
||||
//! - `SetEndOfFile` + `SetFileValidData` for extent preallocation
|
||||
//! (caller needs `SE_MANAGE_VOLUME_NAME` privilege; if unavailable
|
||||
//! fall back to a write-zero path or just skip).
|
||||
//! - `FlushFileBuffers` for fsync-equivalent durable flush.
|
||||
//!
|
||||
//! Until then: preallocate is a debug-logged no-op; durable_sync calls
|
||||
//! the std `File::sync_all` (which on Windows maps to
|
||||
//! `FlushFileBuffers` internally).
|
||||
|
||||
use std::fs::File;
|
||||
use std::io;
|
||||
|
||||
pub(super) fn preallocate(_file: &File, size_bytes: u64) {
|
||||
tracing::debug!(
|
||||
target: "mux",
|
||||
"WritebackFile preallocate size_hint={size_bytes} skipped (windows stub; TODO: SetFileValidData)"
|
||||
);
|
||||
}
|
||||
|
||||
pub(super) fn durable_sync(file: &File) -> io::Result<()> {
|
||||
// `File::sync_all` on Windows is `FlushFileBuffers`. Acceptable
|
||||
// for now; the bounded-syscall wrapper is not used here because
|
||||
// the stub also skips the worker-thread + leak machinery (the
|
||||
// wrapper would need an `unsafe impl Send` for `RawHandle`, and
|
||||
// designing that without a Windows test env is asking for it).
|
||||
file.sync_all()
|
||||
}
|
||||
+4
-2
@@ -154,7 +154,7 @@ pub use profile::DriveProfile;
|
||||
// don't touch `DecryptKeys` directly — `DiscStream::new(reader, title, keys, …)`
|
||||
// accepts whatever `Disc::decrypt_keys()` returned. `decrypt_sectors()` is
|
||||
// for callers that operate on raw sector buffers (e.g. ISO patching).
|
||||
pub use decrypt::{DecryptKeys, decrypt_sectors};
|
||||
pub use decrypt::{DecryptKeys, decrypt_sectors, decrypt_threads, set_decrypt_threads};
|
||||
|
||||
// ─── Disc structure ─────────────────────────────────────────────────────────
|
||||
//
|
||||
@@ -211,9 +211,11 @@ pub use mux::{InputOptions, StreamUrl, input, output, parse_url};
|
||||
// are the ISO-on-disk implementations. [`DecryptingSectorSource`] is the
|
||||
// single decrypt-on-read decorator (AACS / CSS / none) — wrap any
|
||||
// `SectorSource` to get plaintext sectors out.
|
||||
pub use mux::build_iso_pipeline;
|
||||
pub use scsi::{DriveInfo, ScsiSense, ScsiTransport, drive_has_disc, list_drives};
|
||||
pub use sector::{
|
||||
DecryptingSectorSource, FileSectorSink, FileSectorSource, SectorSink, SectorSource,
|
||||
DecryptingSectorSource, FileSectorSink, FileSectorSource, PrefetchedSectorSource, SectorSink,
|
||||
SectorSource,
|
||||
};
|
||||
pub use speed::DriveSpeed;
|
||||
pub use udf::{UdfFs, read_filesystem};
|
||||
|
||||
@@ -76,6 +76,7 @@ impl CodecParser for Ac3Parser {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: data[start..start + frame_size].to_vec(),
|
||||
duration_ns: None,
|
||||
});
|
||||
pos = start + frame_size;
|
||||
}
|
||||
|
||||
@@ -86,6 +86,7 @@ impl CodecParser for DtsParser {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: data[start..start + total_size].to_vec(),
|
||||
duration_ns: None,
|
||||
});
|
||||
pos = start + total_size;
|
||||
}
|
||||
|
||||
@@ -30,6 +30,7 @@ impl CodecParser for DvdSubParser {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
@@ -85,6 +85,7 @@ impl CodecParser for H264Parser {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: frame_data,
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
@@ -176,14 +177,16 @@ impl<'a> Iterator for NalIterator<'a> {
|
||||
}
|
||||
|
||||
/// Find the position of the next start code (00 00 01) at or after `from`.
|
||||
///
|
||||
/// Backed by `memchr::memmem::find` for SIMD-accelerated bytestring
|
||||
/// search. On AVX2-capable x86_64 this runs ~5–10× the byte-by-byte
|
||||
/// scan that preceded it; on a 200 KB UHD HEVC frame the saving is
|
||||
/// in the hundreds of microseconds per call.
|
||||
pub fn find_start_code(data: &[u8], from: usize) -> Option<usize> {
|
||||
if data.len() < from + 3 {
|
||||
return None;
|
||||
}
|
||||
// Range excludes last 2 bytes since we read 3 bytes at each position.
|
||||
// data.len()-2 as exclusive upper bound means last checked index is data.len()-3,
|
||||
// which accesses data[len-3], data[len-2], data[len-1] — all valid.
|
||||
(from..data.len() - 2).find(|&i| data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01)
|
||||
memchr::memmem::find(&data[from..], b"\x00\x00\x01").map(|rel| from + rel)
|
||||
}
|
||||
|
||||
/// Skip past the start code at position `pos`, returning the first byte after it.
|
||||
|
||||
@@ -52,7 +52,10 @@ impl CodecParser for HevcParser {
|
||||
let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
|
||||
let data = &pes.data;
|
||||
let mut keyframe = false;
|
||||
let mut frame_data = Vec::new();
|
||||
// Pre-size: output is ~input bytes with a few 4-byte length
|
||||
// prefixes added. UHD frames are 150-300 KB; the unsized Vec
|
||||
// growth chain otherwise reallocs 5-7× per frame.
|
||||
let mut frame_data = Vec::with_capacity(data.len() + 64);
|
||||
|
||||
// Single-pass NAL scan: extract params, detect keyframes, build length-prefixed output
|
||||
let mut pos = 0;
|
||||
@@ -107,6 +110,7 @@ impl CodecParser for HevcParser {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: frame_data,
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
@@ -44,6 +44,7 @@ impl CodecParser for LpcmParser {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data[BD_LPCM_HEADER_SIZE..].to_vec(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
@@ -29,6 +29,13 @@ pub struct Frame {
|
||||
pub keyframe: bool,
|
||||
/// Frame data (elementary stream bytes).
|
||||
pub data: Vec<u8>,
|
||||
/// Optional duration in nanoseconds — only set by parsers that
|
||||
/// can compute one (currently PGS, which pairs a display PCS
|
||||
/// with the following empty PCS). When `Some`, the MKV muxer
|
||||
/// emits a `BlockGroup` with `BlockDuration` instead of a
|
||||
/// `SimpleBlock`; without it players guess the display interval
|
||||
/// (subtitles linger past their end-time).
|
||||
pub duration_ns: Option<u64>,
|
||||
}
|
||||
|
||||
/// Convert 90kHz PTS to nanoseconds (round to nearest).
|
||||
@@ -71,6 +78,7 @@ impl CodecParser for PassthroughParser {
|
||||
pts_ns,
|
||||
keyframe: self.keyframe,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
@@ -194,6 +194,7 @@ impl CodecParser for Mpeg2Parser {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
+132
-29
@@ -1,12 +1,32 @@
|
||||
//! HDMV PGS (Presentation Graphics Stream) subtitle parser.
|
||||
//!
|
||||
//! PGS segments: PCS, WDS, PDS, ODS, END.
|
||||
//! Each PES packet contains one or more segments.
|
||||
//! All segments are keyframes (no inter-segment dependencies).
|
||||
//! PGS segments: PCS, WDS, PDS, ODS, END. Each PES packet starts with
|
||||
//! one of those (segment_type byte at offset 0).
|
||||
//!
|
||||
//! Subtitle display lifecycle (BD spec):
|
||||
//! - A "display" PCS (number_of_composition_objects > 0) starts a
|
||||
//! visible subtitle. Its WDS/PDS/ODS follow.
|
||||
//! - A later "empty" PCS (number_of_composition_objects == 0) clears
|
||||
//! the screen.
|
||||
//!
|
||||
//! For Matroska output we collapse that pair into one block with
|
||||
//! `BlockDuration` set to (clear_pts - display_pts). Without a
|
||||
//! duration, hardware players linger on the last bitmap until the
|
||||
//! next subtitle replaces it — which can be many seconds, and on a
|
||||
//! disc where the final subtitle has no follower, until end of file.
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
pub struct PgsParser;
|
||||
const SEGMENT_PCS: u8 = 0x16;
|
||||
// Offset within the PES payload at which number_of_composition_objects
|
||||
// lives in a PCS: 3-byte segment header + 10 bytes of PCS fields
|
||||
// (video_w/h, frame_rate, comp_num, comp_state, palette_update,
|
||||
// palette_id_ref) = 13.
|
||||
const PCS_NUM_OBJECTS_OFFSET: usize = 13;
|
||||
|
||||
pub struct PgsParser {
|
||||
pending: Option<(i64, Vec<u8>)>,
|
||||
}
|
||||
|
||||
impl Default for PgsParser {
|
||||
fn default() -> Self {
|
||||
@@ -16,7 +36,7 @@ impl Default for PgsParser {
|
||||
|
||||
impl PgsParser {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
Self { pending: None }
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,11 +46,63 @@ impl CodecParser for PgsParser {
|
||||
return Vec::new();
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
}]
|
||||
|
||||
let is_pcs = pes.data[0] == SEGMENT_PCS;
|
||||
let pcs_num_objects = if is_pcs && pes.data.len() > PCS_NUM_OBJECTS_OFFSET {
|
||||
Some(pes.data[PCS_NUM_OBJECTS_OFFSET])
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
let mut out = Vec::new();
|
||||
match pcs_num_objects {
|
||||
// Clear/empty PCS — closes any pending display. Drop the
|
||||
// clear segment itself; BlockDuration covers the screen
|
||||
// wipe.
|
||||
Some(0) => {
|
||||
if let Some((start_pts, data)) = self.pending.take() {
|
||||
let duration = pts_ns.saturating_sub(start_pts).max(0) as u64;
|
||||
out.push(Frame {
|
||||
pts_ns: start_pts,
|
||||
keyframe: true,
|
||||
data,
|
||||
duration_ns: Some(duration),
|
||||
});
|
||||
}
|
||||
}
|
||||
// Display PCS — start a new pending. If a prior display
|
||||
// was never explicitly cleared (replace-without-clear),
|
||||
// emit it with the new PCS's PTS as its end.
|
||||
Some(_) => {
|
||||
if let Some((start_pts, data)) = self.pending.take() {
|
||||
let duration = pts_ns.saturating_sub(start_pts).max(0) as u64;
|
||||
out.push(Frame {
|
||||
pts_ns: start_pts,
|
||||
keyframe: true,
|
||||
data,
|
||||
duration_ns: Some(duration),
|
||||
});
|
||||
}
|
||||
self.pending = Some((pts_ns, pes.data.clone()));
|
||||
}
|
||||
// Non-PCS first segment — either a continuation of the
|
||||
// current display set, or non-standard layout. If we have
|
||||
// a pending display, append; otherwise emit as-is.
|
||||
None => {
|
||||
if let Some((_, ref mut buf)) = self.pending {
|
||||
buf.extend_from_slice(&pes.data);
|
||||
} else {
|
||||
out.push(Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
duration_ns: None,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
@@ -52,29 +124,60 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_basic_segment() {
|
||||
let mut parser = PgsParser::new();
|
||||
// PGS segment data (PCS = presentation composition segment)
|
||||
let data = vec![0x16, 0x00, 0x00, 0x11, 0x01, 0x02, 0x03];
|
||||
let pes = make_pes(data.clone(), Some(90000));
|
||||
let frames = parser.parse(&pes);
|
||||
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].data, data);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
// Minimum-viable PCS bytes: type 0x16, segment_length (2 bytes),
|
||||
// then 11 bytes of PCS fields ending in number_of_composition_objects.
|
||||
fn pcs_bytes(num_objects: u8) -> Vec<u8> {
|
||||
let mut v = vec![SEGMENT_PCS, 0x00, 0x0B];
|
||||
v.extend_from_slice(&[0x07, 0x80, 0x04, 0x38]); // 1920x1080
|
||||
v.push(0x10); // frame_rate
|
||||
v.extend_from_slice(&[0x00, 0x01]); // composition_number
|
||||
v.push(0x80); // composition_state = EpochStart
|
||||
v.push(0x00); // palette_update + reserved
|
||||
v.push(0x00); // palette_id_ref
|
||||
v.push(num_objects);
|
||||
v
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn all_keyframes() {
|
||||
fn display_then_clear_yields_duration() {
|
||||
let mut parser = PgsParser::new();
|
||||
for i in 0..3 {
|
||||
let data = vec![0x16, 0x00, i];
|
||||
let pes = make_pes(data, Some(90000 * i as i64));
|
||||
let frames = parser.parse(&pes);
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert!(frames[0].keyframe, "PGS segment should always be keyframe");
|
||||
}
|
||||
|
||||
// Display PCS at PTS 90000 (= 1s)
|
||||
let display = pcs_bytes(1);
|
||||
let frames = parser.parse(&make_pes(display.clone(), Some(90000)));
|
||||
assert!(frames.is_empty(), "display PCS should be pending");
|
||||
|
||||
// Empty PCS at PTS 270000 (= 3s)
|
||||
let clear = pcs_bytes(0);
|
||||
let frames = parser.parse(&make_pes(clear, Some(270000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
assert_eq!(frames[0].duration_ns, Some(2_000_000_000));
|
||||
assert_eq!(frames[0].data, display);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn replace_without_clear_still_emits_prior_with_duration() {
|
||||
let mut parser = PgsParser::new();
|
||||
let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000)));
|
||||
let frames = parser.parse(&make_pes(pcs_bytes(1), Some(180000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||||
assert_eq!(frames[0].duration_ns, Some(1_000_000_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_pcs_segment_appends_to_pending() {
|
||||
let mut parser = PgsParser::new();
|
||||
let _ = parser.parse(&make_pes(pcs_bytes(1), Some(90000)));
|
||||
// ODS-like segment (type 0x15)
|
||||
let frames = parser.parse(&make_pes(vec![0x15, 0x00, 0x02, 0xAA, 0xBB], Some(90000)));
|
||||
assert!(frames.is_empty());
|
||||
// Clear closes the set; data should include the appended bytes.
|
||||
let frames = parser.parse(&make_pes(pcs_bytes(0), Some(180000)));
|
||||
assert_eq!(frames.len(), 1);
|
||||
let data = &frames[0].data;
|
||||
assert!(data.windows(5).any(|w| w == [0x15, 0x00, 0x02, 0xAA, 0xBB]));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -139,6 +139,7 @@ impl CodecParser for TrueHdParser {
|
||||
pts_ns: self.next_pts_ns,
|
||||
keyframe: is_major_sync,
|
||||
data: self.buf[..unit_bytes].to_vec(),
|
||||
duration_ns: None,
|
||||
});
|
||||
self.buf.drain(..unit_bytes);
|
||||
self.next_pts_ns += AU_DURATION_NS;
|
||||
|
||||
@@ -96,6 +96,7 @@ impl CodecParser for Vc1Parser {
|
||||
pts_ns: ts_ns,
|
||||
keyframe,
|
||||
data: frame_data.to_vec(),
|
||||
duration_ns: None,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,350 @@
|
||||
//! `DemuxThread` — runs the read+decrypt+demux pipeline on a
|
||||
//! dedicated thread, feeding completed `PesPacket` batches to the
|
||||
//! caller via a bounded channel.
|
||||
//!
|
||||
//! ## Why a second worker thread
|
||||
//!
|
||||
//! With [`crate::sector::PrefetchedSectorSource`] alone, read+decrypt
|
||||
//! already runs on a producer thread; the *consumer* (main) thread
|
||||
//! still serialises `ts_demuxer.feed` (M2TS parsing) with the codec
|
||||
//! parsers. Profiling on the rip1 testbed showed feed at ~37 % and
|
||||
//! codec parse at ~44 % of consumer wall time — i.e. feed is heavy
|
||||
//! enough that pipelining it with parse pays for itself.
|
||||
//!
|
||||
//! Splitting them: feed runs in [`DemuxThread`]; the consumer thread
|
||||
//! receives `Vec<PesPacket>` batches and runs codec parse + frame
|
||||
//! emission only. Total throughput becomes `1/max(feed, parse)`
|
||||
//! instead of `1/(feed + parse)`.
|
||||
//!
|
||||
//! ## Lifecycle
|
||||
//!
|
||||
//! [`DemuxThread::spawn`] takes ownership of the inner reader and the
|
||||
//! demuxer state, returns a handle plus a `Receiver<DemuxBatch>`.
|
||||
//! Dropping the handle closes the channel which signals the thread
|
||||
//! to exit; the join in `Drop::drop` is bounded.
|
||||
|
||||
use crate::halt::Halt;
|
||||
use crate::sector::SectorSource;
|
||||
use crossbeam_channel::{Receiver, Sender, bounded};
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
/// Output channel depth. Two batches in flight keeps the consumer
|
||||
/// (codec parser) busy without piling up demuxed bytes if it stalls.
|
||||
const DEMUX_CHANNEL_DEPTH: usize = 2;
|
||||
|
||||
/// One demuxed batch flowing from the demux thread to the consumer.
|
||||
pub enum DemuxBatch {
|
||||
/// Successfully demuxed PesPackets — non-empty.
|
||||
Ts(Vec<super::ts::PesPacket>),
|
||||
Ps(Vec<super::ps::PsPacket>),
|
||||
/// Underlying reader returned an error. Terminal.
|
||||
Err(std::io::Error),
|
||||
}
|
||||
|
||||
/// Spawned demux thread. Drop joins.
|
||||
///
|
||||
/// In zero-copy mode the thread also owns an opaque
|
||||
/// `producer_shell: Option<Box<dyn Send>>` — the join handle of the
|
||||
/// upstream producer (sector or byte prefetcher). Dropping the
|
||||
/// `DemuxThread` runs the shell's `Drop`, which joins the producer.
|
||||
/// `Box<dyn Send>` rather than a concrete type so the same demux
|
||||
/// worker can be wired behind either prefetcher kind.
|
||||
pub struct DemuxThread {
|
||||
handle: Option<JoinHandle<()>>,
|
||||
#[allow(dead_code)]
|
||||
producer_shell: Option<Box<dyn Send>>,
|
||||
}
|
||||
|
||||
impl DemuxThread {
|
||||
/// Spawn the demux thread. Returns the thread handle and a
|
||||
/// receiver for [`DemuxBatch`] items.
|
||||
///
|
||||
/// `reader` is the fully-composed read+decrypt stack (e.g.
|
||||
/// [`PrefetchedSectorSource`](crate::sector::PrefetchedSectorSource)
|
||||
/// wrapping
|
||||
/// [`DecryptingSectorSource`](crate::sector::DecryptingSectorSource)).
|
||||
/// `extents` is what the thread walks; it issues one
|
||||
/// `read_sectors` per batch of `batch_sectors` sectors (aligned
|
||||
/// to 3-sector AACS units when possible).
|
||||
pub fn spawn<S: SectorSource + Send + 'static>(
|
||||
mut reader: S,
|
||||
extents: Vec<crate::disc::Extent>,
|
||||
batch_sectors: u16,
|
||||
halt: Option<Halt>,
|
||||
ts: Option<super::ts::TsDemuxer>,
|
||||
ps: Option<super::ps::PsDemuxer>,
|
||||
) -> (Self, Receiver<DemuxBatch>) {
|
||||
let (tx, rx) = bounded::<DemuxBatch>(DEMUX_CHANNEL_DEPTH);
|
||||
let mut ts = ts;
|
||||
let mut ps = ps;
|
||||
|
||||
let handle = std::thread::Builder::new()
|
||||
.name("freemkv-demux".into())
|
||||
.spawn(move || {
|
||||
let mut buf = vec![0u8; batch_sectors as usize * 2048];
|
||||
let mut ext_idx = 0usize;
|
||||
let mut offset: u32 = 0;
|
||||
let prof = std::env::var_os("FREEMKV_PROFILE").is_some();
|
||||
let mut prof_started = std::time::Instant::now();
|
||||
let mut prof_last_dump = prof_started;
|
||||
let mut prof_read_ns: u128 = 0;
|
||||
let mut prof_feed_ns: u128 = 0;
|
||||
let mut prof_send_ns: u128 = 0;
|
||||
let mut prof_bytes: u64 = 0;
|
||||
while ext_idx < extents.len() {
|
||||
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
|
||||
return;
|
||||
}
|
||||
let ext = &extents[ext_idx];
|
||||
let remaining = ext.sector_count.saturating_sub(offset);
|
||||
if remaining == 0 {
|
||||
ext_idx += 1;
|
||||
offset = 0;
|
||||
continue;
|
||||
}
|
||||
let mut sectors = remaining.min(batch_sectors as u32) as u16;
|
||||
if sectors >= 3 {
|
||||
sectors -= sectors % 3;
|
||||
}
|
||||
let bytes = sectors as usize * 2048;
|
||||
if buf.len() < bytes {
|
||||
buf.resize(bytes, 0);
|
||||
}
|
||||
let lba = ext.start_lba + offset;
|
||||
let t0 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let n = match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) {
|
||||
Ok(n) => n,
|
||||
Err(e) => {
|
||||
let _ = tx.send(DemuxBatch::Err(e.into()));
|
||||
return;
|
||||
}
|
||||
};
|
||||
let t1 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
offset += sectors as u32;
|
||||
|
||||
// Demux this batch immediately so the channel
|
||||
// carries already-parsed PesPackets, not raw
|
||||
// sector bytes.
|
||||
if let Some(ref mut d) = ts {
|
||||
let pkts = d.feed(&buf[..n]);
|
||||
let t2 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !pkts.is_empty() && tx.send(DemuxBatch::Ts(pkts)).is_err() {
|
||||
return; // consumer dropped
|
||||
}
|
||||
let t3 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if prof {
|
||||
prof_read_ns += t1.unwrap().duration_since(t0.unwrap()).as_nanos();
|
||||
prof_feed_ns += t2.unwrap().duration_since(t1.unwrap()).as_nanos();
|
||||
prof_send_ns += t3.unwrap().duration_since(t2.unwrap()).as_nanos();
|
||||
prof_bytes += n as u64;
|
||||
let now = t3.unwrap();
|
||||
if now.duration_since(prof_last_dump)
|
||||
>= std::time::Duration::from_secs(5)
|
||||
{
|
||||
let el = now.duration_since(prof_started).as_millis().max(1);
|
||||
let mbps = prof_bytes as u128 * 1000 / 1_000_000 / el;
|
||||
eprintln!(
|
||||
"[demux] elapsed={}ms in={}MB/s read={}% feed={}% send={}%",
|
||||
el,
|
||||
mbps,
|
||||
prof_read_ns / 10_000 / el,
|
||||
prof_feed_ns / 10_000 / el,
|
||||
prof_send_ns / 10_000 / el,
|
||||
);
|
||||
prof_last_dump = now;
|
||||
prof_started = now;
|
||||
prof_read_ns = 0;
|
||||
prof_feed_ns = 0;
|
||||
prof_send_ns = 0;
|
||||
prof_bytes = 0;
|
||||
}
|
||||
}
|
||||
} else if let Some(ref mut d) = ps {
|
||||
let pkts = d.feed(&buf[..n]);
|
||||
if !pkts.is_empty() && tx.send(DemuxBatch::Ps(pkts)).is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
// EOF — emit any flushed packets too.
|
||||
if let Some(ref mut d) = ts {
|
||||
let tail = d.flush();
|
||||
if !tail.is_empty() {
|
||||
let _ = tx.send(DemuxBatch::Ts(tail));
|
||||
}
|
||||
} else if let Some(ref mut d) = ps {
|
||||
let tail = d.flush();
|
||||
if !tail.is_empty() {
|
||||
let _ = tx.send(DemuxBatch::Ps(tail));
|
||||
}
|
||||
}
|
||||
// Sender drops here -> consumer sees RecvError → EOF.
|
||||
})
|
||||
.expect("freemkv-demux thread spawn failed");
|
||||
|
||||
(
|
||||
Self {
|
||||
handle: Some(handle),
|
||||
producer_shell: None,
|
||||
},
|
||||
rx,
|
||||
)
|
||||
}
|
||||
|
||||
/// Zero-copy variant. Instead of taking a `SectorSource` and
|
||||
/// memcpy-ing through its `read_sectors` API, this constructor
|
||||
/// consumes the prefetch channels directly: filled buffers come
|
||||
/// in via `prefetch_rx`, the demux thread feeds them, then
|
||||
/// returns them to `recycle_tx` for the producer to re-fill.
|
||||
/// Eliminates the 16 MiB memcpy per batch that the SectorSource
|
||||
/// adapter incurred (and, with the producer-side recycling pool,
|
||||
/// also eliminates the per-batch heap alloc / cross-thread free
|
||||
/// that was costing 40 %+ of demux-thread time before).
|
||||
///
|
||||
/// `producer_shell` is an opaque handle whose only purpose is to
|
||||
/// outlive the demux thread and join the upstream producer when
|
||||
/// dropped. Both
|
||||
/// [`crate::sector::PrefetchedSectorSource::into_channels`] and
|
||||
/// [`crate::io::byte_prefetcher::BytePrefetcher::into_channels`]
|
||||
/// hand back a shell that fits — pass either.
|
||||
pub fn spawn_zero_copy<S: Send + 'static>(
|
||||
prefetch_rx: Receiver<std::io::Result<Vec<u8>>>,
|
||||
recycle_tx: Sender<Vec<u8>>,
|
||||
producer_shell: S,
|
||||
halt: Option<Halt>,
|
||||
ts: Option<super::ts::TsDemuxer>,
|
||||
ps: Option<super::ps::PsDemuxer>,
|
||||
) -> (Self, Receiver<DemuxBatch>) {
|
||||
let (tx, rx) = bounded::<DemuxBatch>(DEMUX_CHANNEL_DEPTH);
|
||||
let mut ts = ts;
|
||||
let mut ps = ps;
|
||||
|
||||
let handle = std::thread::Builder::new()
|
||||
.name("freemkv-demux".into())
|
||||
.spawn(move || {
|
||||
let prof = std::env::var_os("FREEMKV_PROFILE").is_some();
|
||||
let mut prof_started = std::time::Instant::now();
|
||||
let mut prof_last_dump = prof_started;
|
||||
let mut prof_read_ns: u128 = 0;
|
||||
let mut prof_feed_ns: u128 = 0;
|
||||
let mut prof_bytes: u64 = 0;
|
||||
loop {
|
||||
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
|
||||
return;
|
||||
}
|
||||
let t0 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let buf = match prefetch_rx.recv() {
|
||||
Ok(Ok(b)) => b,
|
||||
Ok(Err(e)) => {
|
||||
let _ = tx.send(DemuxBatch::Err(e));
|
||||
return;
|
||||
}
|
||||
Err(_) => break, // producer done → EOF
|
||||
};
|
||||
let t1 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let n = buf.len();
|
||||
if let Some(ref mut d) = ts {
|
||||
let pkts = d.feed(&buf);
|
||||
let t2 = if prof {
|
||||
Some(std::time::Instant::now())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// Recycle the buffer back to the producer
|
||||
// before pushing the demuxed packets. If the
|
||||
// recycle channel is closed the producer has
|
||||
// exited; we drop the buffer and continue.
|
||||
let _ = recycle_tx.send(buf);
|
||||
if !pkts.is_empty() && tx.send(DemuxBatch::Ts(pkts)).is_err() {
|
||||
return;
|
||||
}
|
||||
if prof {
|
||||
prof_read_ns += t1.unwrap().duration_since(t0.unwrap()).as_nanos();
|
||||
prof_feed_ns += t2.unwrap().duration_since(t1.unwrap()).as_nanos();
|
||||
prof_bytes += n as u64;
|
||||
let now = std::time::Instant::now();
|
||||
if now.duration_since(prof_last_dump)
|
||||
>= std::time::Duration::from_secs(5)
|
||||
{
|
||||
let el = now.duration_since(prof_started).as_millis().max(1);
|
||||
let mbps = prof_bytes as u128 * 1000 / 1_000_000 / el;
|
||||
eprintln!(
|
||||
"[demux] elapsed={}ms in={}MB/s read={}% feed={}%",
|
||||
el,
|
||||
mbps,
|
||||
prof_read_ns / 10_000 / el,
|
||||
prof_feed_ns / 10_000 / el,
|
||||
);
|
||||
prof_last_dump = now;
|
||||
prof_started = now;
|
||||
prof_read_ns = 0;
|
||||
prof_feed_ns = 0;
|
||||
prof_bytes = 0;
|
||||
}
|
||||
}
|
||||
} else if let Some(ref mut d) = ps {
|
||||
let pkts = d.feed(&buf);
|
||||
let _ = recycle_tx.send(buf);
|
||||
if !pkts.is_empty() && tx.send(DemuxBatch::Ps(pkts)).is_err() {
|
||||
return;
|
||||
}
|
||||
} else {
|
||||
let _ = recycle_tx.send(buf);
|
||||
}
|
||||
}
|
||||
// Flush tail packets at EOF.
|
||||
if let Some(ref mut d) = ts {
|
||||
let tail = d.flush();
|
||||
if !tail.is_empty() {
|
||||
let _ = tx.send(DemuxBatch::Ts(tail));
|
||||
}
|
||||
} else if let Some(ref mut d) = ps {
|
||||
let tail = d.flush();
|
||||
if !tail.is_empty() {
|
||||
let _ = tx.send(DemuxBatch::Ps(tail));
|
||||
}
|
||||
}
|
||||
})
|
||||
.expect("freemkv-demux thread spawn failed");
|
||||
|
||||
(
|
||||
Self {
|
||||
handle: Some(handle),
|
||||
producer_shell: Some(Box::new(producer_shell)),
|
||||
},
|
||||
rx,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for DemuxThread {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.handle.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
+112
-8
@@ -1,7 +1,7 @@
|
||||
//! DiscStream — read any disc (physical drive or ISO file) → PES frames.
|
||||
//!
|
||||
//! One stream type for all disc sources. The source is a SectorSource —
|
||||
//! Drive (hardware) or IsoSectorReader (file). DiscStream doesn't care.
|
||||
//! Drive (hardware) or FileSectorSource (file). DiscStream doesn't care.
|
||||
//!
|
||||
//! Read-only. For disc→ISO (raw sector copy), use `Disc::copy()`.
|
||||
|
||||
@@ -149,7 +149,11 @@ pub struct DiscStream {
|
||||
// a percent without a separate API call.
|
||||
bytes_total_extents: u64,
|
||||
|
||||
// PES output
|
||||
// PES output — single-threaded inline demux + codec parse. The
|
||||
// pipeline-mode mux (3-stage threaded) lives in
|
||||
// [`super::pipelined_stream::PipelinedPesStream`]; this type is
|
||||
// the legacy in-thread path for live-disc reads where adaptive
|
||||
// batch retry on bad sectors lives in `fill_extents`.
|
||||
ts_demuxer: Option<super::ts::TsDemuxer>,
|
||||
ps_demuxer: Option<super::ps::PsDemuxer>,
|
||||
parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
@@ -403,6 +407,71 @@ impl DiscStream {
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-stage profiling state — populated only when `FREEMKV_PROFILE`
|
||||
/// is set. Dumps a percentage breakdown to stderr every
|
||||
/// [`PROFILE_INTERVAL`]. Zero overhead in normal runs (Option check
|
||||
/// is the only added cost).
|
||||
struct StageProf {
|
||||
started: std::time::Instant,
|
||||
last_dump: std::time::Instant,
|
||||
fill_ns: u128,
|
||||
feed_ns: u128,
|
||||
consume_ns: u128,
|
||||
bytes_in: u64,
|
||||
}
|
||||
|
||||
const PROFILE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5);
|
||||
|
||||
thread_local! {
|
||||
static STAGE_PROF: std::cell::RefCell<Option<StageProf>> = const { std::cell::RefCell::new(None) };
|
||||
}
|
||||
|
||||
fn prof_active() -> bool {
|
||||
std::env::var_os("FREEMKV_PROFILE").is_some()
|
||||
}
|
||||
|
||||
fn prof_tick(stage: &str, ns: u128, bytes: u64) {
|
||||
STAGE_PROF.with(|cell| {
|
||||
let mut slot = cell.borrow_mut();
|
||||
if slot.is_none() {
|
||||
if !prof_active() {
|
||||
return;
|
||||
}
|
||||
let now = std::time::Instant::now();
|
||||
*slot = Some(StageProf {
|
||||
started: now,
|
||||
last_dump: now,
|
||||
fill_ns: 0,
|
||||
feed_ns: 0,
|
||||
consume_ns: 0,
|
||||
bytes_in: 0,
|
||||
});
|
||||
}
|
||||
let p = slot.as_mut().unwrap();
|
||||
match stage {
|
||||
"fill" => p.fill_ns += ns,
|
||||
"feed" => p.feed_ns += ns,
|
||||
"consume" => p.consume_ns += ns,
|
||||
_ => {}
|
||||
}
|
||||
p.bytes_in += bytes;
|
||||
let now = std::time::Instant::now();
|
||||
if now.duration_since(p.last_dump) < PROFILE_INTERVAL {
|
||||
return;
|
||||
}
|
||||
let elapsed_ms = now.duration_since(p.started).as_millis().max(1);
|
||||
let fill_pct = p.fill_ns / 10_000 / elapsed_ms;
|
||||
let feed_pct = p.feed_ns / 10_000 / elapsed_ms;
|
||||
let consume_pct = p.consume_ns / 10_000 / elapsed_ms;
|
||||
let mbps = p.bytes_in as u128 * 1000 / 1_000_000 / elapsed_ms;
|
||||
eprintln!(
|
||||
"[profile] elapsed={}ms in={}MB/s fill={}% feed={}% consume={}%",
|
||||
elapsed_ms, mbps, fill_pct, feed_pct, consume_pct,
|
||||
);
|
||||
p.last_dump = now;
|
||||
});
|
||||
}
|
||||
|
||||
impl crate::pes::Stream for DiscStream {
|
||||
fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
@@ -414,6 +483,7 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
|
||||
loop {
|
||||
let t0 = std::time::Instant::now();
|
||||
if !self.fill_extents()? {
|
||||
self.eof = true;
|
||||
// Flush demuxer — last PES packet may still be in the assembler
|
||||
@@ -475,6 +545,8 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
|
||||
let bytes = self.buf_valid;
|
||||
let t1 = std::time::Instant::now();
|
||||
prof_tick("fill", t1.duration_since(t0).as_nanos(), bytes as u64);
|
||||
// Plaintext: the wrapped reader (DecryptingSectorSource)
|
||||
// applied AACS / CSS in-place during fill_extents'
|
||||
// read_sectors call. The pre-0.18 inline decrypt step
|
||||
@@ -482,21 +554,42 @@ impl crate::pes::Stream for DiscStream {
|
||||
|
||||
if let Some(ref mut demuxer) = self.ts_demuxer {
|
||||
let packets = demuxer.feed(&self.read_buf[..bytes]);
|
||||
for pes in &packets {
|
||||
if let Some((_, track)) =
|
||||
self.pid_to_track.iter().find(|(pid, _)| *pid == pes.pid)
|
||||
let t2 = std::time::Instant::now();
|
||||
prof_tick("feed", t2.duration_since(t1).as_nanos(), 0);
|
||||
let skip_parse = std::env::var_os("FREEMKV_SKIP_PARSE").is_some();
|
||||
for pes in packets {
|
||||
if let Some((_, track)) = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(pid, _)| *pid == pes.pid)
|
||||
.copied()
|
||||
{
|
||||
if let Some((_, parser)) =
|
||||
if skip_parse {
|
||||
// Profiling escape hatch — bypass the codec
|
||||
// parser and pass the raw PES bytes straight
|
||||
// through as a single PesFrame. Lets us
|
||||
// attribute consumer-thread time to
|
||||
// "demux + framing" vs "codec parse".
|
||||
self.pending_frames.push_back(crate::pes::PesFrame {
|
||||
track,
|
||||
pts: pes.pts.map(super::codec::pts_to_ns).unwrap_or(0),
|
||||
keyframe: false,
|
||||
data: pes.data,
|
||||
duration_ns: None,
|
||||
});
|
||||
} else if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(pes) {
|
||||
for frame in parser.parse(&pes) {
|
||||
self.pending_frames.push_back(
|
||||
crate::pes::PesFrame::from_codec_frame(*track, frame),
|
||||
crate::pes::PesFrame::from_codec_frame(track, frame),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
let t3 = std::time::Instant::now();
|
||||
prof_tick("consume", t3.duration_since(t2).as_nanos(), 0);
|
||||
} else if let Some(ref mut demuxer) = self.ps_demuxer {
|
||||
let packets = demuxer.feed(&self.read_buf[..bytes]);
|
||||
for ps in &packets {
|
||||
@@ -570,6 +663,13 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
|
||||
fn headers_ready(&self) -> bool {
|
||||
// FREEMKV_SKIP_PARSE bypasses codec parsers entirely for
|
||||
// bottleneck profiling, so codec_private is never populated.
|
||||
// Pretend headers are ready immediately in that mode so the
|
||||
// CLI loop doesn't hang waiting for them.
|
||||
if std::env::var_os("FREEMKV_SKIP_PARSE").is_some() {
|
||||
return true;
|
||||
}
|
||||
for (idx, s) in self.title.streams.iter().enumerate() {
|
||||
if let crate::disc::Stream::Video(v) = s {
|
||||
if !v.secondary && self.codec_private(idx).is_none() {
|
||||
@@ -579,6 +679,10 @@ impl crate::pes::Stream for DiscStream {
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn errors(&self) -> u64 {
|
||||
self.errors
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -321,6 +321,12 @@ pub const EBML_DOC_TYPE_READ_VERSION: u32 = 0x4285;
|
||||
// Segment
|
||||
pub const SEGMENT: u32 = 0x1853_8067;
|
||||
|
||||
// SeekHead
|
||||
pub const SEEK_HEAD: u32 = 0x114D_9B74;
|
||||
pub const SEEK: u32 = 0x4DBB;
|
||||
pub const SEEK_ID: u32 = 0x53AB;
|
||||
pub const SEEK_POSITION: u32 = 0x53AC;
|
||||
|
||||
// Segment Info
|
||||
pub const INFO: u32 = 0x1549_A966;
|
||||
pub const TIMESTAMP_SCALE: u32 = 0x2A_D7B1;
|
||||
@@ -366,6 +372,9 @@ pub const BIT_DEPTH: u32 = 0x6264;
|
||||
pub const CLUSTER: u32 = 0x1F43_B675;
|
||||
pub const CLUSTER_TIMESTAMP: u32 = 0xE7;
|
||||
pub const SIMPLE_BLOCK: u32 = 0xA3;
|
||||
pub const BLOCK_GROUP: u32 = 0xA0;
|
||||
pub const BLOCK: u32 = 0xA1;
|
||||
pub const BLOCK_DURATION: u32 = 0x9B;
|
||||
|
||||
// Cues
|
||||
pub const CUES: u32 = 0x1C53_BB6B;
|
||||
|
||||
@@ -0,0 +1,404 @@
|
||||
//! Fragmented MP4 muxer — **stub** for Phase 3.
|
||||
//!
|
||||
//! Goal: ISO/IEC 14496-12 fragmented MP4 (`ftyp` + `moov` init segment,
|
||||
//! then a sequence of `moof+mdat` media fragments) targeting a
|
||||
//! [`SequentialSink`](crate::io::sink::SequentialSink). DASH-friendly,
|
||||
//! no Cues backpatch.
|
||||
//!
|
||||
//! Status (v0.21.0 Phase 3): **STUB**. We ship the init segment
|
||||
//! (`ftyp` + a minimal HEVC `moov` skeleton with one video track) so
|
||||
//! the muxer's shape and call site are validated, but media fragments
|
||||
//! are NOT yet emitted — calls to [`Fmp4Mux::write_video`] currently
|
||||
//! accumulate frames into an internal buffer and discard them on
|
||||
//! [`Fmp4Mux::finish`].
|
||||
//!
|
||||
//! ## What's TODO (tracked in Phase 4 / v0.22.0 scope)
|
||||
//!
|
||||
//! - `moof` box: `mfhd` (sequence_number) + `traf` (`tfhd` + `tfdt`
|
||||
//! + `trun` with sample sizes, durations, flags, composition offsets).
|
||||
//! - `mdat` box: concatenated sample data.
|
||||
//! - Fragment cadence: one fragment per GOP or every N seconds,
|
||||
//! whichever comes first.
|
||||
//! - HEVC `hvcC` box inside `moov.trak.mdia.minf.stbl.stsd` so the
|
||||
//! init segment is self-describing.
|
||||
//! - Sample-flags computation (sync vs. delta, depends_on, etc.).
|
||||
//! - Edit lists / fragment_duration for accurate seeking.
|
||||
//!
|
||||
//! Reference: ISO/IEC 14496-12 §8 (Movie Fragments).
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
// Box type literals — four-character codes per ISO/IEC 14496-12 §4.2.
|
||||
|
||||
const FTYP: [u8; 4] = *b"ftyp";
|
||||
const MOOV: [u8; 4] = *b"moov";
|
||||
const MVHD: [u8; 4] = *b"mvhd";
|
||||
const TRAK: [u8; 4] = *b"trak";
|
||||
const TKHD: [u8; 4] = *b"tkhd";
|
||||
const MDIA: [u8; 4] = *b"mdia";
|
||||
const MDHD: [u8; 4] = *b"mdhd";
|
||||
const HDLR: [u8; 4] = *b"hdlr";
|
||||
const MINF: [u8; 4] = *b"minf";
|
||||
const VMHD: [u8; 4] = *b"vmhd";
|
||||
const DINF: [u8; 4] = *b"dinf";
|
||||
const DREF: [u8; 4] = *b"dref";
|
||||
const URL_: [u8; 4] = *b"url ";
|
||||
const STBL: [u8; 4] = *b"stbl";
|
||||
const STSD: [u8; 4] = *b"stsd";
|
||||
const STTS: [u8; 4] = *b"stts";
|
||||
const STSC: [u8; 4] = *b"stsc";
|
||||
const STSZ: [u8; 4] = *b"stsz";
|
||||
const STCO: [u8; 4] = *b"stco";
|
||||
const MVEX: [u8; 4] = *b"mvex";
|
||||
const TREX: [u8; 4] = *b"trex";
|
||||
|
||||
/// Default movie timescale — 90 kHz lines up with MPEG-TS PTS and the
|
||||
/// HEVC SPS `vui_time_scale` for film content, simplifying the math
|
||||
/// when fragment emission lands.
|
||||
const MOVIE_TIMESCALE: u32 = 90_000;
|
||||
/// Video track ID. fMP4 init segments conventionally use track_ID=1
|
||||
/// for the primary video track; a single-track DASH representation has
|
||||
/// no reason to deviate.
|
||||
const VIDEO_TRACK_ID: u32 = 1;
|
||||
|
||||
/// Fragmented MP4 muxer — stub.
|
||||
///
|
||||
/// See the module-level doc comment for what is and isn't shipped in
|
||||
/// this stub.
|
||||
pub struct Fmp4Mux<W: Write> {
|
||||
writer: W,
|
||||
header_written: bool,
|
||||
/// Pending frames — held for the future fragment-emit path. The
|
||||
/// stub drops these on `finish` but keeping them around lets the
|
||||
/// post-stub work re-attach without changing the public API.
|
||||
pending: Vec<PendingSample>,
|
||||
/// hvcC bytes, if provided. Embedded in the `moov.…stsd.hvc1.hvcC`
|
||||
/// box once that path lands.
|
||||
#[allow(dead_code)]
|
||||
codec_private: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
struct PendingSample {
|
||||
#[allow(dead_code)]
|
||||
pts_ns: i64,
|
||||
#[allow(dead_code)]
|
||||
keyframe: bool,
|
||||
#[allow(dead_code)]
|
||||
data: Vec<u8>,
|
||||
}
|
||||
|
||||
impl<W: Write> Fmp4Mux<W> {
|
||||
pub fn new(writer: W) -> Self {
|
||||
Self {
|
||||
writer,
|
||||
header_written: false,
|
||||
pending: Vec::new(),
|
||||
codec_private: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Provide the `HEVCDecoderConfigurationRecord` for the video track.
|
||||
/// The stub stores it but doesn't yet embed it in `moov` — that's
|
||||
/// part of the post-stub work.
|
||||
pub fn set_video_codec_private(&mut self, hvcc: Vec<u8>) {
|
||||
self.codec_private = Some(hvcc);
|
||||
}
|
||||
|
||||
/// Write one video PES frame.
|
||||
///
|
||||
/// **Stub behaviour:** the first call emits the init segment
|
||||
/// (`ftyp` + `moov`) so any consumer that just wants the shape can
|
||||
/// receive it. Subsequent calls accumulate frames in memory for
|
||||
/// the future fragmenting path; **no media bytes are written yet**.
|
||||
pub fn write_video(&mut self, pts_ns: i64, keyframe: bool, data: &[u8]) -> io::Result<()> {
|
||||
if !self.header_written {
|
||||
self.write_init_segment()?;
|
||||
self.header_written = true;
|
||||
}
|
||||
// TODO(0.22.0): emit one `moof+mdat` per GOP. For now stash the
|
||||
// frame so the future patch can hot-wire emission without API
|
||||
// churn.
|
||||
self.pending.push(PendingSample {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: data.to_vec(),
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Flush. The stub additionally drops accumulated `pending` frames.
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
// TODO(0.22.0): emit final fragment from pending; today the
|
||||
// stub just clears the buffer to release memory.
|
||||
self.pending.clear();
|
||||
self.writer.flush()
|
||||
}
|
||||
|
||||
fn write_init_segment(&mut self) -> io::Result<()> {
|
||||
let ftyp = build_ftyp();
|
||||
let moov = build_moov();
|
||||
self.writer.write_all(&ftyp)?;
|
||||
self.writer.write_all(&moov)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Build the `ftyp` box. `major_brand = "iso6"`, `minor_version = 1`,
|
||||
/// compatible brands `iso6 dash msdh hvc1` — the same conservative set
|
||||
/// shaka-packager uses for HEVC-in-fMP4 outputs.
|
||||
fn build_ftyp() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(b"iso6");
|
||||
body.extend_from_slice(&1u32.to_be_bytes());
|
||||
body.extend_from_slice(b"iso6");
|
||||
body.extend_from_slice(b"dash");
|
||||
body.extend_from_slice(b"msdh");
|
||||
body.extend_from_slice(b"hvc1");
|
||||
wrap_box(&FTYP, &body)
|
||||
}
|
||||
|
||||
/// Build the `moov` box — minimal skeleton. Single video trak, no
|
||||
/// hvcC inside stsd yet (TODO: full hvc1 sample entry).
|
||||
fn build_moov() -> Vec<u8> {
|
||||
let mvhd = build_mvhd();
|
||||
let trak = build_video_trak();
|
||||
let mvex = build_mvex();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&mvhd);
|
||||
body.extend_from_slice(&trak);
|
||||
body.extend_from_slice(&mvex);
|
||||
wrap_box(&MOOV, &body)
|
||||
}
|
||||
|
||||
fn build_mvhd() -> Vec<u8> {
|
||||
// Version 0, 100 bytes total body. Fields per ISO/IEC 14496-12 §8.2.2.
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // creation_time
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // modification_time
|
||||
body.extend_from_slice(&MOVIE_TIMESCALE.to_be_bytes());
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // duration = 0 (fragmented)
|
||||
body.extend_from_slice(&0x0001_0000u32.to_be_bytes()); // rate 1.0
|
||||
body.extend_from_slice(&0x0100u16.to_be_bytes()); // volume 1.0
|
||||
body.extend_from_slice(&[0u8; 2]); // reserved
|
||||
body.extend_from_slice(&[0u8; 8]); // reserved
|
||||
// 3x3 identity transformation matrix in 16.16 fixed point.
|
||||
for v in [0x1_0000u32, 0, 0, 0, 0x1_0000, 0, 0, 0, 0x4000_0000] {
|
||||
body.extend_from_slice(&v.to_be_bytes());
|
||||
}
|
||||
body.extend_from_slice(&[0u8; 24]); // pre_defined[6]
|
||||
body.extend_from_slice(&2u32.to_be_bytes()); // next_track_ID (1 reserved for video)
|
||||
wrap_box(&MVHD, &body)
|
||||
}
|
||||
|
||||
fn build_video_trak() -> Vec<u8> {
|
||||
let tkhd = build_tkhd();
|
||||
let mdia = build_mdia();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&tkhd);
|
||||
body.extend_from_slice(&mdia);
|
||||
wrap_box(&TRAK, &body)
|
||||
}
|
||||
|
||||
fn build_tkhd() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
// version=0 | flags=0x000007 (track_enabled | in_movie | in_preview)
|
||||
body.extend_from_slice(&[0, 0, 0, 7]);
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // creation_time
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // modification_time
|
||||
body.extend_from_slice(&VIDEO_TRACK_ID.to_be_bytes());
|
||||
body.extend_from_slice(&[0u8; 4]); // reserved
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // duration
|
||||
body.extend_from_slice(&[0u8; 8]); // reserved
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // layer
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // alternate_group
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // volume (video=0)
|
||||
body.extend_from_slice(&[0u8; 2]); // reserved
|
||||
// 3x3 identity matrix.
|
||||
for v in [0x1_0000u32, 0, 0, 0, 0x1_0000, 0, 0, 0, 0x4000_0000] {
|
||||
body.extend_from_slice(&v.to_be_bytes());
|
||||
}
|
||||
// width / height in 16.16 fixed point — placeholder 1920x1080.
|
||||
body.extend_from_slice(&(1920u32 << 16).to_be_bytes());
|
||||
body.extend_from_slice(&(1080u32 << 16).to_be_bytes());
|
||||
wrap_box(&TKHD, &body)
|
||||
}
|
||||
|
||||
fn build_mdia() -> Vec<u8> {
|
||||
let mdhd = build_mdhd();
|
||||
let hdlr = build_hdlr_vide();
|
||||
let minf = build_minf();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&mdhd);
|
||||
body.extend_from_slice(&hdlr);
|
||||
body.extend_from_slice(&minf);
|
||||
wrap_box(&MDIA, &body)
|
||||
}
|
||||
|
||||
fn build_mdhd() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // creation_time
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // modification_time
|
||||
body.extend_from_slice(&MOVIE_TIMESCALE.to_be_bytes());
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // duration
|
||||
// language: 'und' in 5-bit-per-char ISO 639-2 packed (bit 15 = 0).
|
||||
body.extend_from_slice(&[0x55, 0xC4]);
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // pre_defined
|
||||
wrap_box(&MDHD, &body)
|
||||
}
|
||||
|
||||
fn build_hdlr_vide() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
body.extend_from_slice(&0u32.to_be_bytes()); // pre_defined
|
||||
body.extend_from_slice(b"vide");
|
||||
body.extend_from_slice(&[0u8; 12]); // reserved
|
||||
body.extend_from_slice(b"VideoHandler\0");
|
||||
wrap_box(&HDLR, &body)
|
||||
}
|
||||
|
||||
fn build_minf() -> Vec<u8> {
|
||||
let vmhd = build_vmhd();
|
||||
let dinf = build_dinf();
|
||||
let stbl = build_stbl();
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&vmhd);
|
||||
body.extend_from_slice(&dinf);
|
||||
body.extend_from_slice(&stbl);
|
||||
wrap_box(&MINF, &body)
|
||||
}
|
||||
|
||||
fn build_vmhd() -> Vec<u8> {
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&[0, 0, 0, 1]); // version + flags=1
|
||||
body.extend_from_slice(&0u16.to_be_bytes()); // graphicsmode
|
||||
body.extend_from_slice(&[0u8; 6]); // opcolor
|
||||
wrap_box(&VMHD, &body)
|
||||
}
|
||||
|
||||
fn build_dinf() -> Vec<u8> {
|
||||
let mut dref_body = Vec::new();
|
||||
dref_body.extend_from_slice(&[0, 0, 0, 0]);
|
||||
dref_body.extend_from_slice(&1u32.to_be_bytes()); // entry_count
|
||||
// url with flags=1 (self-contained) and zero name.
|
||||
let url_body = [0u8, 0, 0, 1];
|
||||
dref_body.extend_from_slice(&wrap_box(&URL_, &url_body));
|
||||
let dref = wrap_box(&DREF, &dref_body);
|
||||
wrap_box(&DINF, &dref)
|
||||
}
|
||||
|
||||
fn build_stbl() -> Vec<u8> {
|
||||
// Stub stsd: empty sample description (zero entries). Replace with
|
||||
// hvc1+hvcC once the fragmenting path lands so the init segment is
|
||||
// actually decodable.
|
||||
let mut stsd_body = Vec::new();
|
||||
stsd_body.extend_from_slice(&[0, 0, 0, 0]);
|
||||
stsd_body.extend_from_slice(&0u32.to_be_bytes()); // entry_count
|
||||
let stsd = wrap_box(&STSD, &stsd_body);
|
||||
|
||||
// Empty stts/stsc/stsz/stco — fragmented init has no samples here.
|
||||
let stts = wrap_box(&STTS, &[0, 0, 0, 0, 0, 0, 0, 0]); // version+flags, count=0
|
||||
let stsc = wrap_box(&STSC, &[0, 0, 0, 0, 0, 0, 0, 0]);
|
||||
let stsz = wrap_box(
|
||||
&STSZ,
|
||||
&[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], // version+flags, sample_size=0, count=0
|
||||
);
|
||||
let stco = wrap_box(&STCO, &[0, 0, 0, 0, 0, 0, 0, 0]);
|
||||
|
||||
let mut body = Vec::new();
|
||||
body.extend_from_slice(&stsd);
|
||||
body.extend_from_slice(&stts);
|
||||
body.extend_from_slice(&stsc);
|
||||
body.extend_from_slice(&stsz);
|
||||
body.extend_from_slice(&stco);
|
||||
wrap_box(&STBL, &body)
|
||||
}
|
||||
|
||||
fn build_mvex() -> Vec<u8> {
|
||||
// trex: track_ID=1, default_sample_description_index=1, others=0.
|
||||
let mut trex_body = Vec::new();
|
||||
trex_body.extend_from_slice(&[0, 0, 0, 0]); // version + flags
|
||||
trex_body.extend_from_slice(&VIDEO_TRACK_ID.to_be_bytes());
|
||||
trex_body.extend_from_slice(&1u32.to_be_bytes()); // default_sample_description_index
|
||||
trex_body.extend_from_slice(&0u32.to_be_bytes()); // default_sample_duration
|
||||
trex_body.extend_from_slice(&0u32.to_be_bytes()); // default_sample_size
|
||||
trex_body.extend_from_slice(&0u32.to_be_bytes()); // default_sample_flags
|
||||
let trex = wrap_box(&TREX, &trex_body);
|
||||
wrap_box(&MVEX, &trex)
|
||||
}
|
||||
|
||||
/// Wrap a box body in `[size:u32-BE][type:4]`. Suitable for any body
|
||||
/// that fits in u32; oversized boxes (size > 4 GiB) need the 64-bit
|
||||
/// large-size extension which we don't generate in the stub.
|
||||
fn wrap_box(box_type: &[u8; 4], body: &[u8]) -> Vec<u8> {
|
||||
let size = (body.len() + 8) as u32;
|
||||
let mut out = Vec::with_capacity(body.len() + 8);
|
||||
out.extend_from_slice(&size.to_be_bytes());
|
||||
out.extend_from_slice(box_type);
|
||||
out.extend_from_slice(body);
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Decode the first box's size + type from `buf`.
|
||||
fn read_box_header(buf: &[u8]) -> (u32, [u8; 4]) {
|
||||
let size = u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]);
|
||||
let bt = [buf[4], buf[5], buf[6], buf[7]];
|
||||
(size, bt)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn init_segment_starts_with_ftyp_then_moov() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = Fmp4Mux::new(&mut sink);
|
||||
// Trigger init emission via a single (stubbed) write.
|
||||
mux.write_video(0, true, &[0x00, 0x00, 0x00, 0x01, 0x40])
|
||||
.unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
let (ftyp_size, ftyp_type) = read_box_header(&sink);
|
||||
assert_eq!(&ftyp_type, b"ftyp");
|
||||
assert!(ftyp_size >= 24, "ftyp too small: {ftyp_size}");
|
||||
|
||||
let (moov_size, moov_type) = read_box_header(&sink[ftyp_size as usize..]);
|
||||
assert_eq!(&moov_type, b"moov");
|
||||
assert!(moov_size > 100, "moov skeleton too small: {moov_size}");
|
||||
|
||||
// Stub guarantee: no media bytes after the init segment.
|
||||
let total = ftyp_size as usize + moov_size as usize;
|
||||
assert_eq!(sink.len(), total, "stub leaked media bytes past moov");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn moov_contains_trak_mvex() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = Fmp4Mux::new(&mut sink);
|
||||
mux.write_video(0, true, &[]).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(sink);
|
||||
|
||||
// Re-emit into a fresh buffer for parsing.
|
||||
let mut buf: Vec<u8> = Vec::new();
|
||||
let mut mux2 = Fmp4Mux::new(&mut buf);
|
||||
mux2.write_video(0, true, &[]).unwrap();
|
||||
mux2.finish().unwrap();
|
||||
drop(mux2);
|
||||
|
||||
// Find moov payload start.
|
||||
let (ftyp_size, _) = read_box_header(&buf);
|
||||
let moov_start = ftyp_size as usize;
|
||||
let (moov_size, _) = read_box_header(&buf[moov_start..]);
|
||||
let moov_payload = &buf[moov_start + 8..moov_start + moov_size as usize];
|
||||
|
||||
// Scan for the trak and mvex four-CC anywhere in the moov payload.
|
||||
let has_trak = moov_payload.windows(4).any(|w| w == b"trak");
|
||||
let has_mvex = moov_payload.windows(4).any(|w| w == b"mvex");
|
||||
assert!(has_trak, "moov missing trak");
|
||||
assert!(has_mvex, "moov missing mvex");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
//! HEVC (H.265) elementary stream muxer — Annex B byte stream.
|
||||
//!
|
||||
//! Consumes [`PesFrame`](crate::pes::PesFrame)s for a single video track
|
||||
//! and writes them as a raw `.hevc` / `.h265` Annex B byte stream:
|
||||
//! `00 00 00 01 | NAL_unit | 00 00 00 01 | NAL_unit | …` with no
|
||||
//! container framing.
|
||||
//!
|
||||
//! On the first frame the muxer emits the codec_private's VPS, SPS, PPS
|
||||
//! (parsed from a `HEVCDecoderConfigurationRecord` in
|
||||
//! `length-prefixed-in-hvcC` form), then converts each PES frame's
|
||||
//! length-prefixed NAL units to Annex B and writes them.
|
||||
//!
|
||||
//! Sequential-only — no Cues, no backpatch. Target sink is any
|
||||
//! [`SequentialSink`](crate::io::sink::SequentialSink): file, socket,
|
||||
//! pipe, anything `Write + Send`.
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
/// Annex B 4-byte start code.
|
||||
pub(crate) const START_CODE: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
|
||||
|
||||
/// HEVC NAL unit type bits live in `(byte0 >> 1) & 0x3F` in Annex B.
|
||||
/// We don't filter NAL types here — the muxer is format-only — but we
|
||||
/// keep the constant as documentation of the field layout.
|
||||
#[allow(dead_code)]
|
||||
const HEVC_NAL_TYPE_MASK: u8 = 0x3F;
|
||||
|
||||
/// Streaming HEVC Annex B muxer.
|
||||
///
|
||||
/// One instance per output stream. Tracks whether parameter sets have
|
||||
/// already been emitted so they're written exactly once at the head of
|
||||
/// the stream, mirroring the convention used by `ffmpeg -c:v copy -f
|
||||
/// hevc`.
|
||||
pub struct HevcMux<W: Write> {
|
||||
writer: W,
|
||||
/// `HEVCDecoderConfigurationRecord` payload (hvcC). Parsed lazily
|
||||
/// on the first `write_frame` so callers can set it after
|
||||
/// construction but before the first frame.
|
||||
codec_private: Option<Vec<u8>>,
|
||||
/// Set once VPS/SPS/PPS have been written to the stream. Subsequent
|
||||
/// frames write only their own NAL units.
|
||||
params_written: bool,
|
||||
}
|
||||
|
||||
impl<W: Write> HevcMux<W> {
|
||||
/// Construct over `writer`. The muxer does not impose any extra
|
||||
/// buffering of its own — the sink owns its write buffering policy
|
||||
/// (see [`LocalFileSink`](crate::io::sink::LocalFileSink) and
|
||||
/// [`SocketSink`](crate::io::sink::SocketSink)).
|
||||
pub fn new(writer: W) -> Self {
|
||||
Self {
|
||||
writer,
|
||||
codec_private: None,
|
||||
params_written: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Provide the `HEVCDecoderConfigurationRecord` (hvcC) so the muxer
|
||||
/// can prepend VPS/SPS/PPS Annex B NALs at stream start. Optional —
|
||||
/// if the PES frames already carry inline parameter sets (some
|
||||
/// upstream demuxers do this), skipping this call is fine.
|
||||
pub fn set_codec_private(&mut self, data: Vec<u8>) {
|
||||
self.codec_private = Some(data);
|
||||
}
|
||||
|
||||
/// Write one PES frame (= one access unit) as Annex B NAL units.
|
||||
///
|
||||
/// Input may be either:
|
||||
/// - Length-prefixed: `[u32-BE len][NAL bytes]` repeated. This is
|
||||
/// the form emitted by libfreemkv's HEVC parser (the MKV-native
|
||||
/// layout). Converted to Annex B.
|
||||
/// - Already Annex B: bytes containing `00 00 00 01` start codes
|
||||
/// anywhere in the buffer. Passed through unchanged.
|
||||
///
|
||||
/// `_pts_ns` is accepted for symmetry with other muxers but ignored
|
||||
/// — Annex B has no timing layer.
|
||||
pub fn write_frame(&mut self, _pts_ns: i64, data: &[u8]) -> io::Result<()> {
|
||||
if !self.params_written {
|
||||
if let Some(cp) = &self.codec_private {
|
||||
if let Some(params) = hvcc_to_annex_b(cp) {
|
||||
self.writer.write_all(¶ms)?;
|
||||
}
|
||||
}
|
||||
self.params_written = true;
|
||||
}
|
||||
let annex_b = length_prefixed_to_annex_b(data);
|
||||
self.writer.write_all(&annex_b)
|
||||
}
|
||||
|
||||
/// Flush the underlying writer. No trailer NAL is needed — an Annex
|
||||
/// B stream ends whenever the file/socket ends.
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
self.writer.flush()
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert a `HEVCDecoderConfigurationRecord` (hvcC) into Annex B NAL
|
||||
/// units. Returns `Some(bytes)` if at least one NAL was extracted, else
|
||||
/// `None`.
|
||||
///
|
||||
/// Layout (per ISO/IEC 14496-15 §8.3.3.1.2):
|
||||
/// - 22-byte fixed header
|
||||
/// - byte 22 = `numOfArrays`
|
||||
/// - each array: `array_completeness:1 | reserved:1 | NAL_unit_type:6`,
|
||||
/// `numNalus:u16-BE`, then `numNalus` × `(nalUnitLength:u16-BE +
|
||||
/// NAL bytes)`.
|
||||
///
|
||||
/// We don't filter on NAL type — VPS (32), SPS (33), PPS (34), and any
|
||||
/// SEI arrays included in hvcC all get the same Annex B treatment.
|
||||
fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> {
|
||||
if hvcc.len() < 23 {
|
||||
return None;
|
||||
}
|
||||
let num_arrays = hvcc[22] as usize;
|
||||
let mut out = Vec::new();
|
||||
let mut offset = 23;
|
||||
for _ in 0..num_arrays {
|
||||
if offset + 3 > hvcc.len() {
|
||||
break;
|
||||
}
|
||||
offset += 1; // array_completeness + nal_type byte
|
||||
let num_nalus = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
|
||||
offset += 2;
|
||||
for _ in 0..num_nalus {
|
||||
if offset + 2 > hvcc.len() {
|
||||
break;
|
||||
}
|
||||
let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
|
||||
offset += 2;
|
||||
if offset + nal_len > hvcc.len() {
|
||||
break;
|
||||
}
|
||||
out.extend_from_slice(&START_CODE);
|
||||
out.extend_from_slice(&hvcc[offset..offset + nal_len]);
|
||||
offset += nal_len;
|
||||
}
|
||||
}
|
||||
if out.is_empty() { None } else { Some(out) }
|
||||
}
|
||||
|
||||
/// Convert length-prefixed NAL units (`[u32-BE len][NAL]` repeated) to
|
||||
/// Annex B (`00 00 00 01 [NAL]` repeated).
|
||||
///
|
||||
/// If the input doesn't parse as length-prefixed (no valid lengths
|
||||
/// extracted), it's returned unchanged on the assumption that it's
|
||||
/// already Annex B — some upstream paths (raw HEVC ES from disc) pass
|
||||
/// Annex B straight through the PES layer.
|
||||
pub(crate) fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(data.len() + (data.len() / 32));
|
||||
let mut offset = 0;
|
||||
while offset + 4 <= data.len() {
|
||||
let len = u32::from_be_bytes([
|
||||
data[offset],
|
||||
data[offset + 1],
|
||||
data[offset + 2],
|
||||
data[offset + 3],
|
||||
]) as usize;
|
||||
offset += 4;
|
||||
if offset + len > data.len() {
|
||||
// Mid-NAL truncation — fall through to the pass-through path
|
||||
// rather than emitting a half-NAL.
|
||||
return data.to_vec();
|
||||
}
|
||||
out.extend_from_slice(&START_CODE);
|
||||
out.extend_from_slice(&data[offset..offset + len]);
|
||||
offset += len;
|
||||
}
|
||||
if out.is_empty() && !data.is_empty() {
|
||||
// No length prefixes found — input is likely already Annex B.
|
||||
return data.to_vec();
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn length_prefixed_converts_to_annex_b() {
|
||||
// Two NALs: [3-byte payload AA BB CC] and [2-byte payload DD EE].
|
||||
let mut buf = Vec::new();
|
||||
buf.extend_from_slice(&3u32.to_be_bytes());
|
||||
buf.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
buf.extend_from_slice(&2u32.to_be_bytes());
|
||||
buf.extend_from_slice(&[0xDD, 0xEE]);
|
||||
|
||||
let got = length_prefixed_to_annex_b(&buf);
|
||||
let want = [
|
||||
0x00, 0x00, 0x00, 0x01, 0xAA, 0xBB, 0xCC, // first NAL
|
||||
0x00, 0x00, 0x00, 0x01, 0xDD, 0xEE, // second NAL
|
||||
];
|
||||
assert_eq!(&got[..], &want[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn already_annex_b_passes_through_when_no_lengths_match() {
|
||||
// A buffer < 4 bytes can't parse a length prefix at all →
|
||||
// pass-through path triggers.
|
||||
let raw = [0xAA, 0xBB, 0xCC];
|
||||
let got = length_prefixed_to_annex_b(&raw);
|
||||
assert_eq!(&got[..], &raw[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mid_nal_truncation_returns_original() {
|
||||
// `[u32-BE 100][only 3 bytes]` — length prefix claims 100 bytes
|
||||
// but the input only has 3 after the prefix. We treat that as
|
||||
// malformed and pass the original buffer through so receivers
|
||||
// can attempt their own recovery.
|
||||
let mut raw = Vec::new();
|
||||
raw.extend_from_slice(&100u32.to_be_bytes());
|
||||
raw.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
let got = length_prefixed_to_annex_b(&raw);
|
||||
assert_eq!(&got[..], &raw[..]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hvcc_extracts_vps_sps_pps() {
|
||||
// Build a minimal-but-valid hvcC: 22-byte header, then 3 arrays
|
||||
// (VPS / SPS / PPS), each with 1 NAL of a 4-byte payload that
|
||||
// we can spot in the output.
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(3); // numOfArrays
|
||||
for (nal_type, payload) in [
|
||||
(32u8, [0x40, 0x01, 0x0C, 0x01]),
|
||||
(33, [0x42, 0x01, 0x01, 0x01]),
|
||||
(34, [0x44, 0x01, 0xC1, 0x72]),
|
||||
] {
|
||||
hvcc.push(nal_type & 0x3F);
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
|
||||
hvcc.extend_from_slice(&(payload.len() as u16).to_be_bytes());
|
||||
hvcc.extend_from_slice(&payload);
|
||||
}
|
||||
|
||||
let annex_b = hvcc_to_annex_b(&hvcc).expect("at least one NAL");
|
||||
// Three NALs × (4-byte start + 4-byte payload) = 24 bytes.
|
||||
assert_eq!(annex_b.len(), 24);
|
||||
assert_eq!(&annex_b[..4], &START_CODE);
|
||||
assert_eq!(&annex_b[8..12], &START_CODE);
|
||||
assert_eq!(&annex_b[16..20], &START_CODE);
|
||||
assert_eq!(annex_b[4], 0x40); // VPS first byte
|
||||
assert_eq!(annex_b[12], 0x42); // SPS first byte
|
||||
assert_eq!(annex_b[20], 0x44); // PPS first byte
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mux_writes_params_then_frames() {
|
||||
// Build hvcC with one SPS to verify params-once semantics.
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(1);
|
||||
hvcc.push(33);
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes());
|
||||
hvcc.extend_from_slice(&3u16.to_be_bytes());
|
||||
hvcc.extend_from_slice(&[0x42, 0x01, 0x01]);
|
||||
|
||||
let mut frame_data = Vec::new();
|
||||
frame_data.extend_from_slice(&2u32.to_be_bytes());
|
||||
frame_data.extend_from_slice(&[0xAA, 0xBB]);
|
||||
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = HevcMux::new(&mut sink);
|
||||
mux.set_codec_private(hvcc);
|
||||
mux.write_frame(0, &frame_data).unwrap();
|
||||
// Second frame — no SPS re-emission.
|
||||
mux.write_frame(40_000_000, &frame_data).unwrap();
|
||||
mux.finish().unwrap();
|
||||
|
||||
// SPS NAL (7 bytes) + 2× frame NAL (6 bytes) = 19 bytes.
|
||||
assert_eq!(sink.len(), 7 + 6 + 6);
|
||||
// Start codes at offsets 0 (SPS), 7 (frame1), 13 (frame2).
|
||||
assert_eq!(&sink[0..4], &START_CODE);
|
||||
assert_eq!(&sink[7..11], &START_CODE);
|
||||
assert_eq!(&sink[13..17], &START_CODE);
|
||||
assert_eq!(sink[4], 0x42);
|
||||
assert_eq!(sink[11], 0xAA);
|
||||
assert_eq!(sink[17], 0xAA);
|
||||
}
|
||||
}
|
||||
-101
@@ -1,101 +0,0 @@
|
||||
//! ISO sector reader — file-backed SectorSource for Blu-ray ISO images.
|
||||
//!
|
||||
//! An ISO is a flat image of 2048-byte sectors. Sector N starts at byte offset N * 2048.
|
||||
//! Used by DiscStream::open_iso() and Disc::scan_image().
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
use crate::sector::SectorSource;
|
||||
use std::fs::File;
|
||||
use std::io::{Read, Seek, SeekFrom};
|
||||
use std::path::Path;
|
||||
|
||||
const SECTOR_SIZE: u64 = 2048;
|
||||
|
||||
/// File-backed sector reader for ISO images.
|
||||
pub struct IsoSectorReader {
|
||||
file: File,
|
||||
capacity: u32,
|
||||
}
|
||||
|
||||
impl IsoSectorReader {
|
||||
pub fn open(path: &str) -> std::io::Result<Self> {
|
||||
let file = File::open(Path::new(path))?;
|
||||
let size = file.metadata()?.len();
|
||||
let sectors = size / SECTOR_SIZE;
|
||||
if sectors > u32::MAX as u64 {
|
||||
return Err(crate::error::Error::IsoTooLarge {
|
||||
path: path.to_string(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let capacity = sectors as u32;
|
||||
Ok(Self { file, capacity })
|
||||
}
|
||||
|
||||
pub fn capacity_sectors(&self) -> u32 {
|
||||
self.capacity
|
||||
}
|
||||
}
|
||||
|
||||
impl SectorSource for IsoSectorReader {
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
lba: u32,
|
||||
count: u16,
|
||||
buf: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
let bytes = count as usize * SECTOR_SIZE as usize;
|
||||
self.file
|
||||
.seek(SeekFrom::Start(lba as u64 * SECTOR_SIZE))
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
self.file
|
||||
.read_exact(&mut buf[..bytes])
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn iso_reader_read_sectors() {
|
||||
let mut data = vec![0u8; 4 * SECTOR_SIZE as usize];
|
||||
for i in 0..4u8 {
|
||||
let offset = i as usize * SECTOR_SIZE as usize;
|
||||
data[offset] = i + 1;
|
||||
data[offset + 2047] = i + 100;
|
||||
}
|
||||
|
||||
let dir = std::env::temp_dir().join("freemkv_test_iso_read");
|
||||
std::fs::write(&dir, &data).unwrap();
|
||||
|
||||
let mut reader = IsoSectorReader::open(dir.to_str().unwrap()).unwrap();
|
||||
assert_eq!(reader.capacity_sectors(), 4);
|
||||
|
||||
let mut buf = [0u8; 2048];
|
||||
reader.read_sectors(0, 1, &mut buf, true).unwrap();
|
||||
assert_eq!(buf[0], 1);
|
||||
assert_eq!(buf[2047], 100);
|
||||
|
||||
reader.read_sectors(2, 1, &mut buf, true).unwrap();
|
||||
assert_eq!(buf[0], 3);
|
||||
assert_eq!(buf[2047], 102);
|
||||
|
||||
std::fs::remove_file(&dir).ok();
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn iso_reader_capacity() {
|
||||
let data = vec![0u8; 10 * SECTOR_SIZE as usize];
|
||||
let dir = std::env::temp_dir().join("freemkv_test_iso_cap");
|
||||
std::fs::write(&dir, &data).unwrap();
|
||||
|
||||
let reader = IsoSectorReader::open(dir.to_str().unwrap()).unwrap();
|
||||
assert_eq!(reader.capacity_sectors(), 10);
|
||||
|
||||
std::fs::remove_file(&dir).ok();
|
||||
}
|
||||
}
|
||||
+138
-226
@@ -1,57 +1,21 @@
|
||||
//! M2tsStream — BD transport stream with embedded metadata header.
|
||||
//! M2tsStream — BD transport stream write sink.
|
||||
//!
|
||||
//! Write: prepends FMKV metadata header, then muxes PES frames into BD-TS.
|
||||
//! Read: extracts metadata header (or scans PMT), then demuxes BD-TS into PES frames.
|
||||
//! Write: prepends FMKV metadata header, then muxes PES frames into
|
||||
//! BD-TS. The read direction lives on the pipeline highway —
|
||||
//! `m2ts://` URLs route through
|
||||
//! [`super::resolve::input`] → `build_m2ts_pipeline` →
|
||||
//! [`super::pipelined_stream::PipelinedPesStream`], so this type is
|
||||
//! write-only.
|
||||
|
||||
use super::{meta, ts};
|
||||
use super::meta;
|
||||
use crate::disc::{DiscTitle, Stream as DiscStream};
|
||||
use std::io::{self, Read, Write};
|
||||
use std::io::{self, Write};
|
||||
|
||||
type PesSetup = (
|
||||
Vec<u16>,
|
||||
Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
Vec<(u16, usize)>,
|
||||
);
|
||||
|
||||
/// Size of initial scan buffer for PMT/stream detection.
|
||||
const SCAN_SIZE: usize = 1024 * 1024;
|
||||
|
||||
enum Mode {
|
||||
Write {
|
||||
muxer: super::tsmux::TsMuxer<Box<dyn Write + Send>>,
|
||||
},
|
||||
Read {
|
||||
reader: Box<dyn Read + Send>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Read as many bytes as possible into buf (multiple read calls if needed).
|
||||
/// Bounded by buf.len() — caller controls max bytes read.
|
||||
fn read_fill(r: &mut impl Read, buf: &mut [u8]) -> io::Result<usize> {
|
||||
let mut total = 0;
|
||||
while total < buf.len() {
|
||||
match r.read(&mut buf[total..]) {
|
||||
Ok(0) => break,
|
||||
Ok(n) => total += n,
|
||||
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
Ok(total)
|
||||
}
|
||||
|
||||
/// BD transport stream with embedded metadata.
|
||||
/// BD transport stream write sink with embedded FMKV metadata
|
||||
/// header.
|
||||
pub struct M2tsStream {
|
||||
disc_title: DiscTitle,
|
||||
mode: Mode,
|
||||
// PES support
|
||||
demuxer: Option<ts::TsDemuxer>,
|
||||
parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
pending_frames: std::collections::VecDeque<crate::pes::PesFrame>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
pes_eof: bool,
|
||||
/// Codec private data per stream (from FMKV header).
|
||||
stored_codec_privates: Vec<Option<Vec<u8>>>,
|
||||
muxer: super::tsmux::TsMuxer<Box<dyn Write + Send>>,
|
||||
}
|
||||
|
||||
impl M2tsStream {
|
||||
@@ -81,205 +45,153 @@ impl M2tsStream {
|
||||
}
|
||||
Ok(Self {
|
||||
disc_title: title.clone(),
|
||||
mode: Mode::Write { muxer },
|
||||
demuxer: None,
|
||||
parsers: Vec::new(),
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
pid_to_track: Vec::new(),
|
||||
pes_eof: false,
|
||||
stored_codec_privates: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn setup_pes(streams: &[DiscStream]) -> PesSetup {
|
||||
let mut pids = Vec::new();
|
||||
let mut parsers: Vec<(u16, Box<dyn super::codec::CodecParser>)> = Vec::new();
|
||||
let mut pid_to_track = Vec::new();
|
||||
for (i, s) in streams.iter().enumerate() {
|
||||
let (pid, codec) = match s {
|
||||
DiscStream::Video(v) => (v.pid, v.codec),
|
||||
DiscStream::Audio(a) => (a.pid, a.codec),
|
||||
DiscStream::Subtitle(s) => (s.pid, s.codec),
|
||||
};
|
||||
pids.push(pid);
|
||||
pid_to_track.push((pid, i));
|
||||
parsers.push((pid, super::codec::parser_for_codec(codec, None)));
|
||||
}
|
||||
(pids, parsers, pid_to_track)
|
||||
}
|
||||
|
||||
/// Open an M2TS stream for reading. Takes any Read source — file, pipe, socket.
|
||||
///
|
||||
/// Tries FMKV metadata header first. Falls back to PMT scan of first 1 MB.
|
||||
pub fn open(mut reader: impl Read + Send + 'static) -> io::Result<Self> {
|
||||
// Read first chunk — enough for FMKV header or PMT scan
|
||||
let mut head = vec![0u8; SCAN_SIZE];
|
||||
let head_len = read_fill(&mut reader, &mut head)?;
|
||||
head.truncate(head_len);
|
||||
|
||||
// Try FMKV metadata header from the buffered head
|
||||
let mut cursor = io::Cursor::new(&head);
|
||||
if let Ok(Some(m)) = meta::read_header(&mut cursor) {
|
||||
let header_end = cursor.position() as usize;
|
||||
let title = m.to_title();
|
||||
let (pids, parsers, pid_to_track) = Self::setup_pes(&title.streams);
|
||||
// Chain: remaining head bytes + rest of reader
|
||||
let remaining_head = &head[header_end..];
|
||||
let chain: Box<dyn Read + Send> =
|
||||
Box::new(io::Cursor::new(remaining_head.to_vec()).chain(reader));
|
||||
return Ok(Self {
|
||||
disc_title: title.clone(),
|
||||
mode: Mode::Read { reader: chain },
|
||||
demuxer: if pids.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(ts::TsDemuxer::new(&pids))
|
||||
},
|
||||
parsers,
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
pid_to_track,
|
||||
pes_eof: false,
|
||||
stored_codec_privates: title.codec_privates,
|
||||
});
|
||||
}
|
||||
|
||||
// No FMKV header — scan head for PMT
|
||||
let streams = ts::scan_streams(&head)
|
||||
.ok_or_else(|| -> io::Error { crate::error::Error::NoStreams.into() })?;
|
||||
|
||||
let (pids, parsers, pid_to_track) = Self::setup_pes(&streams);
|
||||
|
||||
// Chain: full head (it's all TS data) + rest of reader
|
||||
let chain: Box<dyn Read + Send> = Box::new(io::Cursor::new(head).chain(reader));
|
||||
|
||||
Ok(Self {
|
||||
disc_title: DiscTitle {
|
||||
duration_secs: 0.0, // unknown without seeking
|
||||
streams,
|
||||
..DiscTitle::empty()
|
||||
},
|
||||
mode: Mode::Read { reader: chain },
|
||||
demuxer: if pids.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(ts::TsDemuxer::new(&pids))
|
||||
},
|
||||
parsers,
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
pid_to_track,
|
||||
pes_eof: false,
|
||||
stored_codec_privates: Vec::new(),
|
||||
muxer,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl crate::pes::Stream for M2tsStream {
|
||||
fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
if self.pes_eof {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
loop {
|
||||
let reader = match &mut self.mode {
|
||||
Mode::Read { reader } => reader,
|
||||
_ => return Err(crate::error::Error::StreamWriteOnly.into()),
|
||||
};
|
||||
let mut buf = vec![0u8; 192 * 1024];
|
||||
let n = reader.read(&mut buf)?;
|
||||
if n == 0 {
|
||||
self.pes_eof = true;
|
||||
// Flush demuxer — last PES packet may still be in the assembler
|
||||
if let Some(ref mut demuxer) = self.demuxer {
|
||||
for pes in &demuxer.flush() {
|
||||
if let Some((_, track)) =
|
||||
self.pid_to_track.iter().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(pes) {
|
||||
self.pending_frames.push_back(
|
||||
crate::pes::PesFrame::from_codec_frame(*track, frame),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return Ok(self.pending_frames.pop_front());
|
||||
}
|
||||
|
||||
if let Some(ref mut demuxer) = self.demuxer {
|
||||
let packets = demuxer.feed(&buf[..n]);
|
||||
for pes in &packets {
|
||||
if let Some((_, track)) =
|
||||
self.pid_to_track.iter().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(pes) {
|
||||
self.pending_frames.push_back(
|
||||
crate::pes::PesFrame::from_codec_frame(*track, frame),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
}
|
||||
// Write-only sink. The m2ts:// read direction is served by
|
||||
// `super::resolve::build_m2ts_pipeline` →
|
||||
// `PipelinedPesStream`; routing through this type for reads
|
||||
// was removed when the highway became the only ingress.
|
||||
Err(crate::error::Error::StreamWriteOnly.into())
|
||||
}
|
||||
|
||||
fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> {
|
||||
match &mut self.mode {
|
||||
Mode::Write { muxer } => muxer.write_frame(frame.track, frame.pts, &frame.data),
|
||||
Mode::Read { .. } => Err(crate::error::Error::StreamReadOnly.into()),
|
||||
}
|
||||
self.muxer
|
||||
.write_frame(frame.track, frame.pts, frame.keyframe, &frame.data)
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> io::Result<()> {
|
||||
match &mut self.mode {
|
||||
Mode::Write { muxer } => muxer.finish(),
|
||||
Mode::Read { .. } => Ok(()),
|
||||
}
|
||||
self.muxer.finish()
|
||||
}
|
||||
|
||||
fn info(&self) -> &crate::disc::DiscTitle {
|
||||
&self.disc_title
|
||||
}
|
||||
|
||||
fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
|
||||
// First check stored codec_privates from FMKV header
|
||||
if let Some(Some(cp)) = self.stored_codec_privates.get(track) {
|
||||
return Some(cp.clone());
|
||||
fn codec_private(&self, _track: usize) -> Option<Vec<u8>> {
|
||||
// Write side doesn't have parsers; codec_private flows in
|
||||
// via the title metadata at `create` time and gets baked
|
||||
// into the FMKV header. Nothing to surface back here.
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::disc::{
|
||||
Codec, ColorSpace, ContentFormat, DiscTitle, FrameRate, HdrFormat, Resolution,
|
||||
Stream as DiscStream, VideoStream,
|
||||
};
|
||||
use crate::pes::{PesFrame, Stream as PesStreamTrait};
|
||||
|
||||
const VIDEO_PID: u16 = 0x1011;
|
||||
|
||||
fn make_title() -> DiscTitle {
|
||||
DiscTitle {
|
||||
playlist: String::new(),
|
||||
playlist_id: 0,
|
||||
duration_secs: 0.0,
|
||||
size_bytes: 0,
|
||||
clips: Vec::new(),
|
||||
streams: vec![DiscStream::Video(VideoStream {
|
||||
pid: VIDEO_PID,
|
||||
codec: Codec::Hevc,
|
||||
resolution: Resolution::R1080p,
|
||||
frame_rate: FrameRate::F24,
|
||||
hdr: HdrFormat::Sdr,
|
||||
color_space: ColorSpace::Bt709,
|
||||
secondary: false,
|
||||
label: String::new(),
|
||||
})],
|
||||
chapters: Vec::new(),
|
||||
extents: Vec::new(),
|
||||
content_format: ContentFormat::BdTs,
|
||||
codec_privates: vec![Some({
|
||||
// Minimal hvcC with one VPS-like array entry.
|
||||
let marker: &[u8] = &[0x40, 0x01, 0x0C, 0x01];
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(1); // numArrays
|
||||
hvcc.push(32);
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
|
||||
hvcc.extend_from_slice(&(marker.len() as u16).to_be_bytes());
|
||||
hvcc.extend_from_slice(marker);
|
||||
hvcc
|
||||
})],
|
||||
}
|
||||
}
|
||||
|
||||
fn fake_idr_pes_data() -> Vec<u8> {
|
||||
// 4-byte length prefix + NAL: type 19 (IDR_W_RADL).
|
||||
let mut nal = vec![(19u8 << 1) & 0x7E, 0x01];
|
||||
for i in 0..200 {
|
||||
nal.push((i & 0xFF) as u8);
|
||||
}
|
||||
// Fall back to parser-extracted codec_private
|
||||
let pid = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(_, idx)| *idx == track)
|
||||
.map(|(pid, _)| *pid)?;
|
||||
self.parsers
|
||||
.iter()
|
||||
.find(|(p, _)| *p == pid)
|
||||
.and_then(|(_, parser)| parser.codec_private())
|
||||
let mut out = Vec::with_capacity(4 + nal.len());
|
||||
out.extend_from_slice(&(nal.len() as u32).to_be_bytes());
|
||||
out.extend_from_slice(&nal);
|
||||
out
|
||||
}
|
||||
|
||||
fn headers_ready(&self) -> bool {
|
||||
for (idx, s) in self.disc_title.streams.iter().enumerate() {
|
||||
if let crate::disc::Stream::Video(v) = s {
|
||||
if !v.secondary && self.codec_private(idx).is_none() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
/// Writer wrapper that shares an Arc<Mutex<Vec<u8>>> so the test can
|
||||
/// inspect the bytes after the muxer drops.
|
||||
struct SharedSink(std::sync::Arc<std::sync::Mutex<Vec<u8>>>);
|
||||
impl Write for SharedSink {
|
||||
fn write(&mut self, b: &[u8]) -> io::Result<usize> {
|
||||
self.0.lock().unwrap().extend_from_slice(b);
|
||||
Ok(b.len())
|
||||
}
|
||||
true
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn m2ts_stream_forwards_keyframe_to_rai() {
|
||||
let title = make_title();
|
||||
let shared = std::sync::Arc::new(std::sync::Mutex::new(Vec::<u8>::new()));
|
||||
let sink = SharedSink(shared.clone());
|
||||
let mut stream = M2tsStream::create(sink, &title).unwrap();
|
||||
let frame = PesFrame {
|
||||
track: 0,
|
||||
pts: 0,
|
||||
keyframe: true,
|
||||
data: fake_idr_pes_data(),
|
||||
duration_ns: None,
|
||||
};
|
||||
stream.write(&frame).unwrap();
|
||||
stream.finish().unwrap();
|
||||
drop(stream);
|
||||
|
||||
let buf = shared.lock().unwrap().clone();
|
||||
|
||||
// Skip FMKV metadata header via meta::read_header.
|
||||
let mut cursor = std::io::Cursor::new(&buf);
|
||||
let _meta = super::meta::read_header(&mut cursor)
|
||||
.unwrap()
|
||||
.expect("FMKV header present");
|
||||
let header_end = cursor.position() as usize;
|
||||
let ts_bytes = &buf[header_end..];
|
||||
|
||||
// Find first PUSI packet on VIDEO_PID; verify RAI in AF flags.
|
||||
let pkt = ts_bytes
|
||||
.chunks(192)
|
||||
.find(|p| {
|
||||
let h = &p[4..];
|
||||
let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16;
|
||||
pid == VIDEO_PID && (h[1] & 0x40) != 0
|
||||
})
|
||||
.expect("video PUSI packet present");
|
||||
let h = &pkt[4..];
|
||||
let afc = (h[3] >> 4) & 0x03;
|
||||
assert!(afc & 0b10 != 0, "AF must be present");
|
||||
let af_len = h[4] as usize;
|
||||
assert!(af_len >= 1, "AF length must include flags byte");
|
||||
let flags = h[5];
|
||||
assert_eq!(flags & 0x40, 0x40, "RAI bit set");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,828 @@
|
||||
//! Standard MPEG-TS (188-byte packets) muxer — sequential-only.
|
||||
//!
|
||||
//! Distinct from `super::tsmux::TsMuxer` (BD-TS with 192-byte packets
|
||||
//! and the 4-byte TP_extra_header). This muxer emits the IETF / ISO/IEC
|
||||
//! 13818-1 wire format that ffmpeg, VLC, and `m2tsindex` consume
|
||||
//! out of the box. Use it for plain `.ts` / `.m2ts` files over a
|
||||
//! [`SequentialSink`](crate::io::sink::SequentialSink), and for
|
||||
//! MPEG-TS-over-UDP via [`UdpSocketSink`](crate::io::sink::UdpSocketSink).
|
||||
//!
|
||||
//! ## Wire format
|
||||
//!
|
||||
//! Every output packet is exactly 188 bytes:
|
||||
//!
|
||||
//! ```text
|
||||
//! sync_byte:8 = 0x47
|
||||
//! transport_error_indicator:1 = 0
|
||||
//! payload_unit_start_indicator:1
|
||||
//! transport_priority:1 = 0
|
||||
//! PID:13
|
||||
//! transport_scrambling_control:2 = 0
|
||||
//! adaptation_field_control:2
|
||||
//! continuity_counter:4
|
||||
//! [adaptation field, if signalled]
|
||||
//! [payload, if signalled]
|
||||
//! ```
|
||||
//!
|
||||
//! ## Wiring
|
||||
//!
|
||||
//! Single program (PMT PID `0x1000`, program number `1`):
|
||||
//! - PAT on PID `0x0000`
|
||||
//! - PMT on PID `0x1000`
|
||||
//! - Video (HEVC, stream_type `0x24`) on PID `0x0100`
|
||||
//! - First audio (AC3, stream_type `0x81`) on PID `0x0101`
|
||||
//! (or TrueHD, stream_type `0x83`, if hinted)
|
||||
//!
|
||||
//! ## Scope vs. full TS
|
||||
//!
|
||||
//! This is a deliberately minimal viable muxer:
|
||||
//! - One program, one video track, optionally one audio track.
|
||||
//! - PAT + PMT re-emitted every `PSI_INTERVAL_PACKETS` packets so a
|
||||
//! mid-stream receiver can lock on within ~100 ms at typical
|
||||
//! UHD bitrates.
|
||||
//! - PCR clock derived from video PTS (`pts - PCR_LEAD_90KHZ`),
|
||||
//! attached to the video PID's adaptation field every
|
||||
//! `PCR_INTERVAL_PACKETS` packets.
|
||||
//! - No language / descriptor tags, no SCTE-35 markers, no per-PID
|
||||
//! PMT version bumps, no SDT/EIT. Sufficient for "ffmpeg can play
|
||||
//! this back", not for full broadcast deployment.
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
mod packet;
|
||||
|
||||
use packet::{Packet, PacketWriter};
|
||||
|
||||
/// PID `0x0000` — PAT, mandated by spec.
|
||||
const PID_PAT: u16 = 0x0000;
|
||||
/// PID for the single program's PMT. `0x1000` is the conventional
|
||||
/// choice; anything outside reserved/null ranges works.
|
||||
const PID_PMT: u16 = 0x1000;
|
||||
/// PID for the video elementary stream.
|
||||
const PID_VIDEO: u16 = 0x0100;
|
||||
/// PID for the (optional) first audio elementary stream.
|
||||
const PID_AUDIO: u16 = 0x0101;
|
||||
/// Null PID — reserved by spec; never used here.
|
||||
#[allow(dead_code)]
|
||||
const PID_NULL: u16 = 0x1FFF;
|
||||
|
||||
/// Re-emit PAT/PMT every N TS packets. ~250 × 188 B = 47 KB; at
|
||||
/// 80 Mb/s UHD that's ~5 ms — well under typical receiver lock budget.
|
||||
const PSI_INTERVAL_PACKETS: u64 = 250;
|
||||
/// Re-stamp PCR every N TS packets carrying the video PID. Spec MAX
|
||||
/// is 100 ms; 40 packets × 188 B at moderate bitrate keeps us under.
|
||||
const PCR_INTERVAL_PACKETS: u64 = 40;
|
||||
/// PCR lead time — the PCR must precede the PTS of the first byte of
|
||||
/// the picture it timestamps. 200 ms in 90 kHz ticks.
|
||||
const PCR_LEAD_90KHZ: u64 = 90_000 / 5;
|
||||
|
||||
/// Stream-type codes from ISO/IEC 13818-1 Table 2-29 + later amendments.
|
||||
const STREAM_TYPE_HEVC: u8 = 0x24;
|
||||
const STREAM_TYPE_AC3: u8 = 0x81;
|
||||
const STREAM_TYPE_TRUEHD: u8 = 0x83;
|
||||
|
||||
/// Audio codec hint for [`M2tsMux::new`] / [`M2tsMux::set_audio`]. The
|
||||
/// muxer needs to know the codec to pick the right PMT `stream_type`
|
||||
/// and the right PES stream_id; it doesn't decode the audio.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum AudioCodec {
|
||||
/// AC-3 / E-AC-3, stream_type `0x81`, PES stream_id `0xBD`.
|
||||
Ac3,
|
||||
/// Dolby TrueHD, stream_type `0x83`, PES stream_id `0xBD`.
|
||||
TrueHd,
|
||||
}
|
||||
|
||||
impl AudioCodec {
|
||||
fn stream_type(self) -> u8 {
|
||||
match self {
|
||||
AudioCodec::Ac3 => STREAM_TYPE_AC3,
|
||||
AudioCodec::TrueHd => STREAM_TYPE_TRUEHD,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Sequential MPEG-TS muxer with one video and optional one audio track.
|
||||
pub struct M2tsMux<W: Write> {
|
||||
out: PacketWriter<W>,
|
||||
/// hvcC parameter set bytes to prepend to the first video PES.
|
||||
/// Optional — if the upstream frames already carry inline params,
|
||||
/// callers can omit this.
|
||||
video_codec_private: Option<Vec<u8>>,
|
||||
/// Set on first video frame: have we emitted VPS/SPS/PPS?
|
||||
params_written: bool,
|
||||
/// Audio codec, if an audio track is configured. `None` = video-only.
|
||||
audio: Option<AudioCodec>,
|
||||
/// First seen PTS (90 kHz). All subsequent PTS / PCR values are
|
||||
/// relative to this so output streams start near t=0 and don't
|
||||
/// confuse downstream parsers that don't tolerate huge starting
|
||||
/// timestamps.
|
||||
base_pts_90k: Option<u64>,
|
||||
/// Per-PID continuity counter, 4 bits, monotonically increasing.
|
||||
cc_video: u8,
|
||||
cc_audio: u8,
|
||||
cc_pat: u8,
|
||||
cc_pmt: u8,
|
||||
/// Total packets written, used to gate PSI / PCR cadence.
|
||||
packets_written: u64,
|
||||
/// Video packets written since last PCR, used to gate PCR cadence.
|
||||
video_packets_since_pcr: u64,
|
||||
}
|
||||
|
||||
impl<W: Write> M2tsMux<W> {
|
||||
/// Construct a muxer wrapping `writer`. By default audio is
|
||||
/// disabled; call [`set_audio`](Self::set_audio) before the first
|
||||
/// frame to enable.
|
||||
pub fn new(writer: W) -> Self {
|
||||
Self {
|
||||
out: PacketWriter::new(writer),
|
||||
video_codec_private: None,
|
||||
params_written: false,
|
||||
audio: None,
|
||||
base_pts_90k: None,
|
||||
cc_video: 0,
|
||||
cc_audio: 0,
|
||||
cc_pat: 0,
|
||||
cc_pmt: 0,
|
||||
packets_written: 0,
|
||||
video_packets_since_pcr: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Provide the `HEVCDecoderConfigurationRecord` for video so the
|
||||
/// muxer prepends VPS/SPS/PPS Annex B NALs at stream start.
|
||||
pub fn set_video_codec_private(&mut self, hvcc: Vec<u8>) {
|
||||
self.video_codec_private = Some(hvcc);
|
||||
}
|
||||
|
||||
/// Enable a single audio track. Must be called before
|
||||
/// [`write_audio`](Self::write_audio).
|
||||
pub fn set_audio(&mut self, codec: AudioCodec) {
|
||||
self.audio = Some(codec);
|
||||
}
|
||||
|
||||
/// Write one video PES frame. `data` is either length-prefixed
|
||||
/// NALUs (MKV-style) or already Annex B; both are accepted. `keyframe`
|
||||
/// drives the random_access_indicator bit on the first packet of this
|
||||
/// PES (and gates codec_private NAL prepending — those only attach to
|
||||
/// the first keyframe).
|
||||
pub fn write_video(&mut self, pts_ns: i64, keyframe: bool, data: &[u8]) -> io::Result<()> {
|
||||
let pts_90k = self.base_relative_pts(pts_ns);
|
||||
// PCR comes "before" the PTS it timestamps; clamp at 0 for the
|
||||
// first frame so we don't underflow.
|
||||
let pcr = pts_90k.saturating_sub(PCR_LEAD_90KHZ);
|
||||
|
||||
// Annex-B-ify the frame and prepend VPS/SPS/PPS once, on the
|
||||
// FIRST keyframe (not first frame — non-key frames before the
|
||||
// first keyframe can't carry params usefully).
|
||||
let mut es = Vec::with_capacity(data.len() + 64);
|
||||
if keyframe && !self.params_written {
|
||||
if let Some(cp) = &self.video_codec_private {
|
||||
let payload = hvcc_payload(cp);
|
||||
if !payload.is_empty() {
|
||||
let params = super::hevc::length_prefixed_to_annex_b(&payload);
|
||||
es.extend_from_slice(¶ms);
|
||||
}
|
||||
}
|
||||
self.params_written = true;
|
||||
}
|
||||
let annex_b = super::hevc::length_prefixed_to_annex_b(data);
|
||||
es.extend_from_slice(&annex_b);
|
||||
|
||||
let pes = build_video_pes(pts_90k, &es);
|
||||
self.write_pes(PID_VIDEO, &pes, Some(pcr), keyframe)
|
||||
}
|
||||
|
||||
/// Write one audio PES frame. Returns `Ok(())` and silently drops
|
||||
/// the frame if no audio track was configured — the design assumes
|
||||
/// the upstream picks tracks and won't ship audio to a video-only
|
||||
/// muxer, but defending against it keeps the API a single shape.
|
||||
pub fn write_audio(&mut self, pts_ns: i64, data: &[u8]) -> io::Result<()> {
|
||||
if self.audio.is_none() {
|
||||
return Ok(());
|
||||
}
|
||||
let pts_90k = self.base_relative_pts(pts_ns);
|
||||
let pes = build_audio_pes(pts_90k, data);
|
||||
self.write_pes(PID_AUDIO, &pes, None, false)
|
||||
}
|
||||
|
||||
/// Drain the underlying writer. No TS-level trailer is mandatory —
|
||||
/// receivers detect end-of-stream from socket close or file EOF.
|
||||
pub fn finish(&mut self) -> io::Result<()> {
|
||||
self.out.flush()
|
||||
}
|
||||
|
||||
/// Convert input PTS (nanoseconds) to 90 kHz ticks rebased on the
|
||||
/// first frame's PTS. Saturating at 0 keeps the math friendly when
|
||||
/// frames arrive slightly out of decode order.
|
||||
fn base_relative_pts(&mut self, pts_ns: i64) -> u64 {
|
||||
let raw_90k = if pts_ns > 0 {
|
||||
(pts_ns as u64) * 9 / 100_000
|
||||
} else {
|
||||
0
|
||||
};
|
||||
let base = *self.base_pts_90k.get_or_insert(raw_90k);
|
||||
raw_90k.saturating_sub(base)
|
||||
}
|
||||
|
||||
/// Emit one PES payload as a chain of TS packets on `pid`. If `pcr`
|
||||
/// is provided the first packet carries an adaptation field with
|
||||
/// the PCR. PAT/PMT are re-emitted every `PSI_INTERVAL_PACKETS`.
|
||||
///
|
||||
/// Packet-size math (TS = 188 bytes, header = 4 bytes):
|
||||
/// - 184 B remain after the header for the adaptation-field area
|
||||
/// plus the payload area.
|
||||
/// - With AF body of `b` bytes and `s` stuffing bytes: AF total =
|
||||
/// `1 + b + s` (the leading `1` is the `adaptation_field_length`
|
||||
/// byte itself). Payload = `184 - (1 + b + s)`.
|
||||
/// - With no AF area at all: payload = `184`.
|
||||
///
|
||||
/// The fit-the-tail logic on the last packet of the PES uses
|
||||
/// stuffing rather than a separate small packet, which is the
|
||||
/// standard MPEG-TS convention.
|
||||
fn write_pes(
|
||||
&mut self,
|
||||
pid: u16,
|
||||
pes: &[u8],
|
||||
pcr: Option<u64>,
|
||||
is_keyframe_video: bool,
|
||||
) -> io::Result<()> {
|
||||
// PSI cadence is enforced per TS packet — interleave a fresh
|
||||
// PAT+PMT into the packet stream every PSI_INTERVAL_PACKETS so
|
||||
// long single-PES emissions (e.g. one 60 KB video frame) don't
|
||||
// starve receivers tuning in mid-stream.
|
||||
let mut offset = 0;
|
||||
let mut first = true;
|
||||
while offset < pes.len() {
|
||||
self.maybe_emit_psi()?;
|
||||
|
||||
let attach_pcr = first
|
||||
&& (pid == PID_VIDEO)
|
||||
&& (pcr.is_some())
|
||||
&& (self.packets_written == 0
|
||||
|| self.video_packets_since_pcr >= PCR_INTERVAL_PACKETS);
|
||||
|
||||
// RAI rides only the FIRST packet of a keyframe video PES.
|
||||
let attach_rai = first && is_keyframe_video && pid == PID_VIDEO;
|
||||
|
||||
let mut af_body: Vec<u8> = if attach_pcr {
|
||||
build_pcr_adaptation(pcr.unwrap_or(0))
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
if attach_rai {
|
||||
if af_body.is_empty() {
|
||||
af_body.push(0x40); // flags: RAI only
|
||||
} else {
|
||||
af_body[0] |= 0x40; // OR RAI into existing PCR flags
|
||||
}
|
||||
}
|
||||
|
||||
let remaining = pes.len() - offset;
|
||||
// Capacity for payload given AF body and 1-byte AF length.
|
||||
// When AF body is empty we can still skip the AF entirely
|
||||
// and get the full 184 B; only invoke the AF when we'd
|
||||
// otherwise need stuffing.
|
||||
let (af_present, payload_len, stuffing): (bool, usize, usize) = if !af_body.is_empty() {
|
||||
// AF is mandatory (PCR). 1 byte length + body + stuffing
|
||||
// + payload = 184.
|
||||
let max_payload = 184 - 1 - af_body.len();
|
||||
let p = remaining.min(max_payload);
|
||||
let s = max_payload - p;
|
||||
(true, p, s)
|
||||
} else if remaining >= 184 {
|
||||
// Full payload packet — no AF at all.
|
||||
(false, 184, 0)
|
||||
} else {
|
||||
// Last (small) packet — stuff via empty AF.
|
||||
// 1 byte length + 0 body + stuffing + payload = 184.
|
||||
let max_payload = 183;
|
||||
let p = remaining.min(max_payload);
|
||||
let s = max_payload - p;
|
||||
(true, p, s)
|
||||
};
|
||||
|
||||
let cc = self.advance_cc(pid);
|
||||
let mut packet = Packet::new();
|
||||
packet.set_header(pid, first, true, af_present, cc);
|
||||
if af_present {
|
||||
packet.append_adaptation(&af_body, stuffing);
|
||||
}
|
||||
packet.append_payload(&pes[offset..offset + payload_len]);
|
||||
debug_assert_eq!(packet.len(), 188, "packet not 188 bytes");
|
||||
self.out.write_packet(&packet)?;
|
||||
|
||||
self.packets_written += 1;
|
||||
if pid == PID_VIDEO {
|
||||
if attach_pcr {
|
||||
self.video_packets_since_pcr = 0;
|
||||
} else {
|
||||
self.video_packets_since_pcr += 1;
|
||||
}
|
||||
}
|
||||
offset += payload_len;
|
||||
first = false;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn advance_cc(&mut self, pid: u16) -> u8 {
|
||||
let slot = match pid {
|
||||
PID_VIDEO => &mut self.cc_video,
|
||||
PID_AUDIO => &mut self.cc_audio,
|
||||
PID_PAT => &mut self.cc_pat,
|
||||
PID_PMT => &mut self.cc_pmt,
|
||||
_ => return 0,
|
||||
};
|
||||
let cc = *slot;
|
||||
*slot = (*slot + 1) & 0x0F;
|
||||
cc
|
||||
}
|
||||
|
||||
fn maybe_emit_psi(&mut self) -> io::Result<()> {
|
||||
if self.packets_written == 0 || self.packets_written % PSI_INTERVAL_PACKETS == 0 {
|
||||
self.emit_pat()?;
|
||||
self.emit_pmt()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn emit_pat(&mut self) -> io::Result<()> {
|
||||
let payload = build_pat(PID_PMT);
|
||||
let cc = self.advance_cc(PID_PAT);
|
||||
let mut packet = Packet::new();
|
||||
packet.set_header(PID_PAT, true, true, false, cc);
|
||||
packet.append_payload(&payload);
|
||||
packet.pad_to_188();
|
||||
self.out.write_packet(&packet)?;
|
||||
self.packets_written += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn emit_pmt(&mut self) -> io::Result<()> {
|
||||
let payload = build_pmt(self.audio);
|
||||
let cc = self.advance_cc(PID_PMT);
|
||||
let mut packet = Packet::new();
|
||||
packet.set_header(PID_PMT, true, true, false, cc);
|
||||
packet.append_payload(&payload);
|
||||
packet.pad_to_188();
|
||||
self.out.write_packet(&packet)?;
|
||||
self.packets_written += 1;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract the raw hvcC bytes for handoff to `length_prefixed_to_annex_b`.
|
||||
/// hvcC layout: 22-byte fixed header, then `numOfArrays` arrays of
|
||||
/// `(nalType, numNalus, [nalLength:u16, NAL bytes]…)`. We convert this
|
||||
/// directly to a length-prefixed byte stream (NAL length is u16 in
|
||||
/// hvcC; widen to u32 for the standard length-prefixed encoding).
|
||||
fn hvcc_payload(hvcc: &[u8]) -> Vec<u8> {
|
||||
if hvcc.len() < 23 {
|
||||
return Vec::new();
|
||||
}
|
||||
let num_arrays = hvcc[22] as usize;
|
||||
let mut out = Vec::new();
|
||||
let mut offset = 23;
|
||||
for _ in 0..num_arrays {
|
||||
if offset + 3 > hvcc.len() {
|
||||
break;
|
||||
}
|
||||
offset += 1;
|
||||
let num_nalus = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
|
||||
offset += 2;
|
||||
for _ in 0..num_nalus {
|
||||
if offset + 2 > hvcc.len() {
|
||||
break;
|
||||
}
|
||||
let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
|
||||
offset += 2;
|
||||
if offset + nal_len > hvcc.len() {
|
||||
break;
|
||||
}
|
||||
out.extend_from_slice(&(nal_len as u32).to_be_bytes());
|
||||
out.extend_from_slice(&hvcc[offset..offset + nal_len]);
|
||||
offset += nal_len;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Build a PES packet for a video access unit.
|
||||
fn build_video_pes(pts_90k: u64, es: &[u8]) -> Vec<u8> {
|
||||
build_pes_packet(0xE0, pts_90k, es, /* length_in_header */ false)
|
||||
}
|
||||
|
||||
/// Build a PES packet for an audio access unit.
|
||||
fn build_audio_pes(pts_90k: u64, es: &[u8]) -> Vec<u8> {
|
||||
// Audio PES: length is fillable when it fits in u16. We always
|
||||
// write the length so receivers don't have to scan for the next
|
||||
// start code.
|
||||
build_pes_packet(0xBD, pts_90k, es, /* length_in_header */ true)
|
||||
}
|
||||
|
||||
fn build_pes_packet(stream_id: u8, pts_90k: u64, es: &[u8], length_in_header: bool) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(es.len() + 14);
|
||||
out.extend_from_slice(&[0x00, 0x00, 0x01, stream_id]);
|
||||
// PES_packet_length: total bytes after this field. 3 flag bytes + 5
|
||||
// PTS bytes + es.len(). Zero means "unbounded" — used for video
|
||||
// where PES can exceed u16.
|
||||
let pes_len = 8 + es.len();
|
||||
if length_in_header && pes_len <= u16::MAX as usize {
|
||||
out.extend_from_slice(&(pes_len as u16).to_be_bytes());
|
||||
} else {
|
||||
out.extend_from_slice(&[0x00, 0x00]);
|
||||
}
|
||||
// Flags byte 1: 10 = MPEG-2 marker, then scrambling/priority/etc 0.
|
||||
out.push(0x80);
|
||||
// Flags byte 2: PTS flag (bit 7).
|
||||
out.push(0x80);
|
||||
// PES_header_data_length = 5 (just PTS).
|
||||
out.push(5);
|
||||
// PTS bytes — 33-bit timestamp split across 5 bytes with marker bits.
|
||||
let pts = pts_90k & 0x1_FFFF_FFFF;
|
||||
out.push(0x21 | (((pts >> 29) & 0x0E) as u8));
|
||||
out.push(((pts >> 22) & 0xFF) as u8);
|
||||
out.push(0x01 | (((pts >> 14) & 0xFE) as u8));
|
||||
out.push(((pts >> 7) & 0xFF) as u8);
|
||||
out.push(0x01 | (((pts << 1) & 0xFE) as u8));
|
||||
out.extend_from_slice(es);
|
||||
out
|
||||
}
|
||||
|
||||
/// Build the PAT payload (section, with pointer_field).
|
||||
fn build_pat(pmt_pid: u16) -> Vec<u8> {
|
||||
let mut section = Vec::new();
|
||||
section.push(0x00); // table_id = PAT
|
||||
// section_syntax_indicator(1) | '0'(1) | reserved(2) | section_length(12)
|
||||
// section_length covers from end of this field through CRC.
|
||||
// Body: transport_stream_id(2) + version/cni(1) + section/last_section(2) + program(4) = 9 bytes,
|
||||
// plus CRC(4) = 13. Encoded big-endian.
|
||||
section.extend_from_slice(&[0xB0, 13]);
|
||||
section.extend_from_slice(&[0x00, 0x01]); // transport_stream_id = 1
|
||||
section.push(0xC1); // reserved | version=0 | current_next=1
|
||||
section.push(0x00); // section_number
|
||||
section.push(0x00); // last_section_number
|
||||
section.extend_from_slice(&[0x00, 0x01]); // program_number = 1
|
||||
// reserved(3) | network_PID/program_map_PID(13)
|
||||
let pid_bytes = (0xE000u16 | (pmt_pid & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&pid_bytes);
|
||||
let crc = mpegts_crc32(§ion);
|
||||
section.extend_from_slice(&crc.to_be_bytes());
|
||||
// Prepend pointer_field=0 (section starts immediately).
|
||||
let mut payload = Vec::with_capacity(section.len() + 1);
|
||||
payload.push(0x00);
|
||||
payload.extend_from_slice(§ion);
|
||||
payload
|
||||
}
|
||||
|
||||
/// Build the PMT payload (section, with pointer_field).
|
||||
fn build_pmt(audio: Option<AudioCodec>) -> Vec<u8> {
|
||||
let mut section = Vec::new();
|
||||
section.push(0x02); // table_id = PMT
|
||||
// section_length filled in after we know the body size.
|
||||
let len_placeholder = section.len();
|
||||
section.extend_from_slice(&[0xB0, 0x00]);
|
||||
|
||||
section.extend_from_slice(&1u16.to_be_bytes()); // program_number
|
||||
section.push(0xC1); // reserved | version=0 | current_next=1
|
||||
section.push(0x00); // section_number
|
||||
section.push(0x00); // last_section_number
|
||||
// reserved(3) | PCR_PID(13)
|
||||
let pcr_pid = (0xE000u16 | (PID_VIDEO & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&pcr_pid);
|
||||
// program_info_length = 0
|
||||
section.extend_from_slice(&[0xF0, 0x00]);
|
||||
|
||||
// Video elementary stream entry.
|
||||
section.push(STREAM_TYPE_HEVC);
|
||||
let v_pid = (0xE000u16 | (PID_VIDEO & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&v_pid);
|
||||
section.extend_from_slice(&[0xF0, 0x00]); // ES_info_length = 0
|
||||
|
||||
if let Some(codec) = audio {
|
||||
section.push(codec.stream_type());
|
||||
let a_pid = (0xE000u16 | (PID_AUDIO & 0x1FFF)).to_be_bytes();
|
||||
section.extend_from_slice(&a_pid);
|
||||
section.extend_from_slice(&[0xF0, 0x00]);
|
||||
}
|
||||
|
||||
// Now patch section_length: covers everything after the length field
|
||||
// through the CRC, so (current body size - 3 bytes consumed by
|
||||
// table_id + 2 length bytes) + 4 (CRC).
|
||||
let section_len = section.len() - 3 + 4;
|
||||
section[len_placeholder] = 0xB0 | ((section_len >> 8) as u8 & 0x0F);
|
||||
section[len_placeholder + 1] = section_len as u8;
|
||||
|
||||
let crc = mpegts_crc32(§ion);
|
||||
section.extend_from_slice(&crc.to_be_bytes());
|
||||
|
||||
let mut payload = Vec::with_capacity(section.len() + 1);
|
||||
payload.push(0x00);
|
||||
payload.extend_from_slice(§ion);
|
||||
payload
|
||||
}
|
||||
|
||||
/// Build the adaptation field carrying a PCR (no other flags).
|
||||
fn build_pcr_adaptation(pcr_90k: u64) -> Vec<u8> {
|
||||
// adaptation_field_length is set by `Packet::append_adaptation`
|
||||
// — this function returns just the field body.
|
||||
//
|
||||
// Layout: discontinuity_indicator(1) | random_access(1) |
|
||||
// elementary_stream_priority(1) | PCR_flag(1) | OPCR_flag(1) |
|
||||
// splicing_point_flag(1) | transport_private_data_flag(1) |
|
||||
// adaptation_field_extension_flag(1) | PCR(48b).
|
||||
let mut af = vec![0x10]; // PCR_flag=1; RAI is OR'd in by the caller when applicable
|
||||
let pcr_base = pcr_90k & 0x1_FFFF_FFFF; // 33-bit
|
||||
let pcr_ext: u16 = 0; // 9-bit, we keep it zero (no sub-tick precision)
|
||||
// Encode PCR: 33b base | 6b reserved | 9b extension = 48b
|
||||
af.push((pcr_base >> 25) as u8);
|
||||
af.push((pcr_base >> 17) as u8);
|
||||
af.push((pcr_base >> 9) as u8);
|
||||
af.push((pcr_base >> 1) as u8);
|
||||
af.push(((pcr_base << 7) as u8 & 0x80) | 0x7E | ((pcr_ext >> 8) as u8 & 0x01));
|
||||
af.push(pcr_ext as u8);
|
||||
af
|
||||
}
|
||||
|
||||
/// MPEG-TS CRC-32 (poly 0x04C11DB7, init 0xFFFFFFFF, no reflection, no
|
||||
/// final XOR). Implementation: bitwise so we don't need a table —
|
||||
/// PSI sections are tiny, the cost is negligible.
|
||||
fn mpegts_crc32(data: &[u8]) -> u32 {
|
||||
let mut crc: u32 = 0xFFFF_FFFF;
|
||||
for &b in data {
|
||||
crc ^= (b as u32) << 24;
|
||||
for _ in 0..8 {
|
||||
if crc & 0x8000_0000 != 0 {
|
||||
crc = (crc << 1) ^ 0x04C1_1DB7;
|
||||
} else {
|
||||
crc <<= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
crc
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// All emitted bytes must align to 188-byte packet boundaries and
|
||||
/// every packet must start with `0x47`.
|
||||
fn assert_ts_well_formed(buf: &[u8]) {
|
||||
assert_eq!(
|
||||
buf.len() % 188,
|
||||
0,
|
||||
"stream not packet-aligned: {} bytes",
|
||||
buf.len()
|
||||
);
|
||||
for (i, chunk) in buf.chunks(188).enumerate() {
|
||||
assert_eq!(chunk[0], 0x47, "packet {} missing sync byte", i);
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_pids(buf: &[u8]) -> Vec<u16> {
|
||||
buf.chunks(188)
|
||||
.map(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crc32_is_self_validating() {
|
||||
// The MPEG-TS CRC has the property that prepending a single
|
||||
// bit-flip changes the output; running it over its own input +
|
||||
// CRC yields a fixed magic constant (the CRC residue). Rather
|
||||
// than hardcoding sample bytes, verify the underlying algorithm
|
||||
// by checking that two distinct inputs produce distinct CRCs
|
||||
// and that the same input is deterministic.
|
||||
let a = [
|
||||
0u8, 0xB0, 0x0D, 0x00, 0x01, 0xC1, 0x00, 0x00, 0x00, 0x01, 0xE1, 0x00,
|
||||
];
|
||||
let mut b = a;
|
||||
b[5] ^= 0x01; // flip one bit
|
||||
let crc_a = mpegts_crc32(&a);
|
||||
let crc_b = mpegts_crc32(&b);
|
||||
assert_ne!(crc_a, crc_b);
|
||||
assert_eq!(crc_a, mpegts_crc32(&a)); // deterministic
|
||||
// Sanity: all-zero input ⇒ CRC = 0 (init XORs but the
|
||||
// shift/feedback cancels for zero data after init drains).
|
||||
// We don't assert exact value — that depends on poly choice —
|
||||
// but check it's not the same as for non-zero data.
|
||||
let crc_zero = mpegts_crc32(&[0u8; 12]);
|
||||
assert_ne!(crc_zero, crc_a);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn video_only_mux_emits_pat_pmt_then_video() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// One small video frame, no codec_private (so no params NAL inline).
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&4u32.to_be_bytes());
|
||||
frame.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]);
|
||||
mux.write_video(0, true, &frame).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
assert_ts_well_formed(&sink);
|
||||
let pids = extract_pids(&sink);
|
||||
// First two packets: PAT, PMT. At least one video packet after.
|
||||
assert_eq!(pids[0], PID_PAT);
|
||||
assert_eq!(pids[1], PID_PMT);
|
||||
assert!(pids.iter().any(|p| *p == PID_VIDEO));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn audio_track_appears_in_pmt_and_stream() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
mux.set_audio(AudioCodec::Ac3);
|
||||
// Video + audio frame pair.
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&3u32.to_be_bytes());
|
||||
frame.extend_from_slice(&[0x40, 0x01, 0x0C]);
|
||||
mux.write_video(0, true, &frame).unwrap();
|
||||
mux.write_audio(20_000_000, &[0x0B, 0x77, 0x12, 0x34])
|
||||
.unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
assert_ts_well_formed(&sink);
|
||||
let pids = extract_pids(&sink);
|
||||
assert!(pids.iter().any(|p| *p == PID_VIDEO));
|
||||
assert!(pids.iter().any(|p| *p == PID_AUDIO));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn psi_re_emits_at_interval() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// Build a frame large enough to span > PSI_INTERVAL_PACKETS TS packets.
|
||||
// 184 B payload per packet ⇒ ~250 packets = ~46 KB elementary stream.
|
||||
let big: Vec<u8> = (0..(60 * 1024)).map(|i| (i & 0xff) as u8).collect();
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&(big.len() as u32).to_be_bytes());
|
||||
frame.extend_from_slice(&big);
|
||||
mux.write_video(0, true, &frame).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
assert_ts_well_formed(&sink);
|
||||
let pids = extract_pids(&sink);
|
||||
// Count PAT/PMT pairs — must be at least 2 given the input size.
|
||||
let pat_count = pids.iter().filter(|p| **p == PID_PAT).count();
|
||||
let pmt_count = pids.iter().filter(|p| **p == PID_PMT).count();
|
||||
assert!(pat_count >= 2, "expected ≥2 PAT, got {}", pat_count);
|
||||
assert!(pmt_count >= 2, "expected ≥2 PMT, got {}", pmt_count);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn continuity_counter_increments_per_pid() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// Three small video frames to get a sequence of video TS packets.
|
||||
for pts in [0i64, 40_000_000, 80_000_000] {
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&3u32.to_be_bytes());
|
||||
frame.extend_from_slice(&[0xAA, 0xBB, 0xCC]);
|
||||
mux.write_video(pts, pts == 0, &frame).unwrap();
|
||||
}
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// Collect CCs for video packets in order.
|
||||
let ccs: Vec<u8> = sink
|
||||
.chunks(188)
|
||||
.filter(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO)
|
||||
.map(|p| p[3] & 0x0F)
|
||||
.collect();
|
||||
for w in ccs.windows(2) {
|
||||
assert_eq!(w[1], (w[0] + 1) & 0x0F);
|
||||
}
|
||||
}
|
||||
|
||||
/// Return the adaptation field body (length byte stripped) for one
|
||||
/// 188-byte TS packet, or None when the packet has no AF.
|
||||
fn af_body(packet: &[u8]) -> Option<Vec<u8>> {
|
||||
let afc = (packet[3] >> 4) & 0x03;
|
||||
if afc & 0b10 == 0 {
|
||||
return None;
|
||||
}
|
||||
let af_len = packet[4] as usize;
|
||||
if af_len == 0 {
|
||||
return Some(Vec::new());
|
||||
}
|
||||
Some(packet[5..5 + af_len].to_vec())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rai_set_on_keyframe_pes_packet() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&4u32.to_be_bytes());
|
||||
frame.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]);
|
||||
mux.write_video(0, true, &frame).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// Find the first PUSI packet on PID_VIDEO.
|
||||
let pkt = sink
|
||||
.chunks(188)
|
||||
.find(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO && (p[1] & 0x40) != 0)
|
||||
.expect("video PUSI packet exists");
|
||||
let af = af_body(pkt).expect("AF present on first packet of keyframe video PES");
|
||||
assert!(!af.is_empty(), "AF flags byte present");
|
||||
assert_eq!(af[0] & 0x40, 0x40, "RAI bit set");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pcr_packet_without_keyframe_has_rai_clear() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
// First frame is the keyframe (gates codec_private; also gets PCR).
|
||||
let mut frame0 = Vec::new();
|
||||
frame0.extend_from_slice(&4u32.to_be_bytes());
|
||||
frame0.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]);
|
||||
mux.write_video(0, true, &frame0).unwrap();
|
||||
// Push enough non-key video frames to cross PCR_INTERVAL_PACKETS
|
||||
// video packets so a later PCR-bearing packet exists.
|
||||
// Each frame is ~50 KB → ~270 packets, well over 40.
|
||||
let big: Vec<u8> = (0..(50 * 1024)).map(|i| (i & 0xff) as u8).collect();
|
||||
for i in 1..3 {
|
||||
let mut frame = Vec::new();
|
||||
frame.extend_from_slice(&(big.len() as u32).to_be_bytes());
|
||||
frame.extend_from_slice(&big);
|
||||
mux.write_video((i as i64) * 40_000_000, false, &frame)
|
||||
.unwrap();
|
||||
}
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// The first PUSI video packet carries PCR + RAI (keyframe).
|
||||
// A later video PUSI packet with AF + PCR but NOT keyframe must
|
||||
// have RAI clear.
|
||||
let video_pusi: Vec<&[u8]> = sink
|
||||
.chunks(188)
|
||||
.filter(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO && (p[1] & 0x40) != 0)
|
||||
.collect();
|
||||
assert!(
|
||||
video_pusi.len() >= 2,
|
||||
"expected ≥2 video PES starts, got {}",
|
||||
video_pusi.len()
|
||||
);
|
||||
// Find a later one with AF that carries PCR (flags & 0x10 set).
|
||||
let later_pcr = video_pusi
|
||||
.iter()
|
||||
.skip(1)
|
||||
.find_map(|p| {
|
||||
let af = af_body(p)?;
|
||||
if !af.is_empty() && (af[0] & 0x10) != 0 {
|
||||
Some(af)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.expect("later PCR-bearing PUSI exists");
|
||||
assert_eq!(
|
||||
later_pcr[0] & 0x40,
|
||||
0,
|
||||
"RAI must be clear on non-keyframe PCR packet"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keyframe_video_with_pcr_combines_flags() {
|
||||
// PCR attaches only when video_packets_since_pcr >=
|
||||
// PCR_INTERVAL_PACKETS (40). The very first video packet emits a
|
||||
// PAT+PMT first, so packets_written != 0 and attach_pcr is false on
|
||||
// frame 0. We push: keyframe (no PCR) → many non-key (drives the
|
||||
// PCR counter past the interval) → second keyframe (PCR + RAI).
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
let mut mux = M2tsMux::new(&mut sink);
|
||||
let mut small = Vec::new();
|
||||
small.extend_from_slice(&4u32.to_be_bytes());
|
||||
small.extend_from_slice(&[0x40, 0x01, 0x0C, 0x01]);
|
||||
mux.write_video(0, true, &small).unwrap();
|
||||
// ~50 KB ≈ 270 packets — well over PCR_INTERVAL_PACKETS.
|
||||
let big: Vec<u8> = (0..(50 * 1024)).map(|i| (i & 0xff) as u8).collect();
|
||||
let mut big_frame = Vec::new();
|
||||
big_frame.extend_from_slice(&(big.len() as u32).to_be_bytes());
|
||||
big_frame.extend_from_slice(&big);
|
||||
mux.write_video(40_000_000, false, &big_frame).unwrap();
|
||||
// Now a second keyframe — must combine RAI (keyframe) and PCR
|
||||
// (counter exceeded).
|
||||
mux.write_video(80_000_000, true, &small).unwrap();
|
||||
mux.finish().unwrap();
|
||||
drop(mux);
|
||||
|
||||
// Collect video PUSI packets and find the third (second keyframe).
|
||||
let video_pusi: Vec<&[u8]> = sink
|
||||
.chunks(188)
|
||||
.filter(|p| u16::from_be_bytes([p[1] & 0x1F, p[2]]) == PID_VIDEO && (p[1] & 0x40) != 0)
|
||||
.collect();
|
||||
assert!(video_pusi.len() >= 3, "three video PES starts expected");
|
||||
let af = af_body(video_pusi[2]).expect("AF present");
|
||||
assert!(!af.is_empty(), "AF flags byte present");
|
||||
assert_eq!(af[0], 0x50, "flags == RAI | PCR");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
//! 188-byte MPEG-TS packet builder + writer.
|
||||
//!
|
||||
//! Internal helper for `super::M2tsMux`. Not public API — exposes raw
|
||||
//! byte layout so the parent module can compose PSI / PCR / PES bytes
|
||||
//! without each caller re-implementing the 188-byte boundary math.
|
||||
|
||||
use std::io::{self, Write};
|
||||
|
||||
const TS_PACKET_SIZE: usize = 188;
|
||||
const SYNC_BYTE: u8 = 0x47;
|
||||
const STUFF_BYTE: u8 = 0xFF;
|
||||
|
||||
/// One TS packet under construction. Always emits 188 bytes when
|
||||
/// [`pad_to_188`](Self::pad_to_188) is called; if it's not called the
|
||||
/// caller is responsible for filling the packet exactly.
|
||||
pub(super) struct Packet {
|
||||
buf: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Packet {
|
||||
pub(super) fn new() -> Self {
|
||||
Self {
|
||||
buf: Vec::with_capacity(TS_PACKET_SIZE),
|
||||
}
|
||||
}
|
||||
|
||||
/// Write the 4-byte TS packet header.
|
||||
///
|
||||
/// * `pid` — 13-bit PID
|
||||
/// * `payload_unit_start` — first packet of a PES / PSI section
|
||||
/// * `has_payload` — packet carries any payload bytes
|
||||
/// * `has_adaptation` — packet carries an adaptation field
|
||||
/// * `cc` — 4-bit continuity counter
|
||||
pub(super) fn set_header(
|
||||
&mut self,
|
||||
pid: u16,
|
||||
payload_unit_start: bool,
|
||||
has_payload: bool,
|
||||
has_adaptation: bool,
|
||||
cc: u8,
|
||||
) {
|
||||
self.buf.clear();
|
||||
self.buf.push(SYNC_BYTE);
|
||||
let pus_bit = if payload_unit_start { 0x40 } else { 0 };
|
||||
// transport_error_indicator(1)=0 | payload_unit_start(1) | transport_priority(1)=0 | PID(5 high)
|
||||
self.buf.push(pus_bit | ((pid >> 8) as u8 & 0x1F));
|
||||
self.buf.push(pid as u8);
|
||||
// transport_scrambling_control(2)=0 | adaptation_field_control(2) | continuity_counter(4)
|
||||
let afc = match (has_adaptation, has_payload) {
|
||||
(false, false) => 0b00, // reserved — should not happen
|
||||
(false, true) => 0b01, // payload only
|
||||
(true, false) => 0b10, // adaptation only
|
||||
(true, true) => 0b11, // both
|
||||
};
|
||||
self.buf.push((afc << 4) | (cc & 0x0F));
|
||||
}
|
||||
|
||||
/// Append the adaptation field after the header.
|
||||
///
|
||||
/// `body` is the adaptation field body (flags byte plus optional PCR
|
||||
/// and so on). `stuffing` is the number of `0xFF` stuffing bytes to
|
||||
/// append after the body. The first byte of the field
|
||||
/// (`adaptation_field_length`) is computed here from
|
||||
/// `body.len() + stuffing`.
|
||||
pub(super) fn append_adaptation(&mut self, body: &[u8], stuffing: usize) {
|
||||
let af_len = body.len() + stuffing;
|
||||
debug_assert!(af_len <= 183, "adaptation field overflow");
|
||||
self.buf.push(af_len as u8);
|
||||
self.buf.extend_from_slice(body);
|
||||
for _ in 0..stuffing {
|
||||
self.buf.push(STUFF_BYTE);
|
||||
}
|
||||
}
|
||||
|
||||
/// Append payload bytes.
|
||||
pub(super) fn append_payload(&mut self, payload: &[u8]) {
|
||||
self.buf.extend_from_slice(payload);
|
||||
debug_assert!(self.buf.len() <= TS_PACKET_SIZE, "packet overflow");
|
||||
}
|
||||
|
||||
/// Pad the packet to exactly 188 bytes with `0xFF` bytes — used by
|
||||
/// PSI emit paths where the section is much smaller than 184 bytes.
|
||||
/// For PSI packets only — payload-carrying packets reserve room for
|
||||
/// stuffing via `append_adaptation`.
|
||||
pub(super) fn pad_to_188(&mut self) {
|
||||
while self.buf.len() < TS_PACKET_SIZE {
|
||||
self.buf.push(STUFF_BYTE);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn bytes(&self) -> &[u8] {
|
||||
&self.buf
|
||||
}
|
||||
|
||||
pub(super) fn len(&self) -> usize {
|
||||
self.buf.len()
|
||||
}
|
||||
}
|
||||
|
||||
/// Buffered writer for assembled TS packets. Owns the underlying sink.
|
||||
pub(super) struct PacketWriter<W: Write> {
|
||||
inner: W,
|
||||
}
|
||||
|
||||
impl<W: Write> PacketWriter<W> {
|
||||
pub(super) fn new(inner: W) -> Self {
|
||||
Self { inner }
|
||||
}
|
||||
|
||||
pub(super) fn write_packet(&mut self, packet: &Packet) -> io::Result<()> {
|
||||
let bytes = packet.bytes();
|
||||
debug_assert_eq!(bytes.len(), TS_PACKET_SIZE);
|
||||
self.inner.write_all(bytes)
|
||||
}
|
||||
|
||||
pub(super) fn flush(&mut self) -> io::Result<()> {
|
||||
self.inner.flush()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn pad_fills_to_188() {
|
||||
let mut p = Packet::new();
|
||||
p.set_header(0x100, true, true, false, 0);
|
||||
p.append_payload(&[1, 2, 3]);
|
||||
p.pad_to_188();
|
||||
assert_eq!(p.bytes().len(), 188);
|
||||
assert_eq!(p.bytes()[0], SYNC_BYTE);
|
||||
// After 4-byte header + 3 payload, byte 7 starts stuffing.
|
||||
assert_eq!(p.bytes()[7], STUFF_BYTE);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn header_pid_round_trips() {
|
||||
let mut p = Packet::new();
|
||||
p.set_header(0x1ABC, false, true, false, 0xA);
|
||||
let pid = u16::from_be_bytes([p.bytes()[1] & 0x1F, p.bytes()[2]]);
|
||||
assert_eq!(pid, 0x1ABC);
|
||||
assert_eq!(p.bytes()[3] & 0x0F, 0xA);
|
||||
}
|
||||
}
|
||||
+564
-25
@@ -87,11 +87,19 @@ impl MkvTrack {
|
||||
}
|
||||
|
||||
pub fn audio(a: &AudioStream) -> Self {
|
||||
// Codec ID strings must distinguish the DTS family — strict
|
||||
// players (Plex transcoder, some hardware decoders, some AV
|
||||
// receivers) reject lossless DTS-HD MA payload when the
|
||||
// track advertises plain `A_DTS` because it implies the
|
||||
// bitstream is the 1.5 Mbps "core" only. Fix is just to
|
||||
// emit the right ID per BD-STN codec field.
|
||||
let codec_id = match a.codec {
|
||||
Codec::Ac3 => "A_AC3",
|
||||
Codec::Ac3Plus => "A_EAC3",
|
||||
Codec::TrueHd => "A_TRUEHD",
|
||||
Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => "A_DTS",
|
||||
Codec::DtsHdMa => "A_DTS/MA",
|
||||
Codec::DtsHdHr => "A_DTS/HR",
|
||||
Codec::Dts => "A_DTS",
|
||||
Codec::Lpcm => "A_PCM/INT/BIG",
|
||||
_ => "A_AC3",
|
||||
};
|
||||
@@ -159,6 +167,12 @@ struct CuePoint {
|
||||
cluster_pos: u64, // relative to Segment start
|
||||
}
|
||||
|
||||
/// SeekHead entry that needs its 8-byte SeekPosition back-patched after Cues are written.
|
||||
struct SeekPositionFixup {
|
||||
target_id: u32,
|
||||
value_offset: u64, // absolute file offset of the 8-byte SeekPosition value
|
||||
}
|
||||
|
||||
/// MKV muxer. Call write_frame() for each frame, then finish() at the end.
|
||||
pub struct MkvMuxer<W: Write + Seek> {
|
||||
writer: W,
|
||||
@@ -170,6 +184,10 @@ pub struct MkvMuxer<W: Write + Seek> {
|
||||
base_pts_ms: Option<i64>,
|
||||
cues: Vec<CuePoint>,
|
||||
frame_count: u64,
|
||||
seek_fixups: Vec<SeekPositionFixup>,
|
||||
info_offset: u64,
|
||||
tracks_offset: u64,
|
||||
chapters_offset: Option<u64>,
|
||||
}
|
||||
|
||||
/// New cluster every 5 seconds.
|
||||
@@ -200,7 +218,34 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::write_unknown_size(&mut writer)?;
|
||||
let segment_start = writer.stream_position()?;
|
||||
|
||||
// SeekHead with fixed-width SeekPosition placeholders. Order: Info, Tracks, [Chapters], Cues.
|
||||
let mut seek_fixups: Vec<SeekPositionFixup> = Vec::new();
|
||||
let seekhead_pos = ebml::start_master(&mut writer, ebml::SEEK_HEAD)?;
|
||||
let mut targets: Vec<u32> = vec![ebml::INFO, ebml::TRACKS];
|
||||
if !chapters.is_empty() {
|
||||
targets.push(ebml::CHAPTERS);
|
||||
}
|
||||
targets.push(ebml::CUES);
|
||||
let seek_id_be = (ebml::SEEK as u16).to_be_bytes();
|
||||
let seek_inner_id_be = (ebml::SEEK_ID as u16).to_be_bytes();
|
||||
let seek_pos_id_be = (ebml::SEEK_POSITION as u16).to_be_bytes();
|
||||
for target_id in &targets {
|
||||
writer.write_all(&[seek_id_be[0], seek_id_be[1], 0x92])?;
|
||||
writer.write_all(&[seek_inner_id_be[0], seek_inner_id_be[1], 0x84])?;
|
||||
writer.write_all(&target_id.to_be_bytes())?;
|
||||
writer.write_all(&[seek_pos_id_be[0], seek_pos_id_be[1], 0x88])?;
|
||||
let value_offset = writer.stream_position()?;
|
||||
writer.write_all(&[0u8; 8])?;
|
||||
seek_fixups.push(SeekPositionFixup {
|
||||
target_id: *target_id,
|
||||
value_offset,
|
||||
});
|
||||
}
|
||||
ebml::end_master(&mut writer, seekhead_pos)?;
|
||||
|
||||
// Info
|
||||
let info_start = writer.stream_position()?;
|
||||
let info_offset = info_start - segment_start;
|
||||
let info_pos = ebml::start_master(&mut writer, ebml::INFO)?;
|
||||
ebml::write_uint(&mut writer, ebml::TIMESTAMP_SCALE, 1_000_000)?; // 1ms precision
|
||||
if duration_secs > 0.0 {
|
||||
@@ -215,6 +260,8 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut writer, info_pos)?;
|
||||
|
||||
// Tracks
|
||||
let tracks_start = writer.stream_position()?;
|
||||
let tracks_offset = tracks_start - segment_start;
|
||||
let tracks_pos = ebml::start_master(&mut writer, ebml::TRACKS)?;
|
||||
for (i, track) in tracks.iter().enumerate() {
|
||||
let entry_pos = ebml::start_master(&mut writer, ebml::TRACK_ENTRY)?;
|
||||
@@ -302,7 +349,10 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut writer, tracks_pos)?;
|
||||
|
||||
// Chapters
|
||||
let mut chapters_offset: Option<u64> = None;
|
||||
if !chapters.is_empty() {
|
||||
let chapters_start = writer.stream_position()?;
|
||||
chapters_offset = Some(chapters_start - segment_start);
|
||||
let chapters_pos = ebml::start_master(&mut writer, ebml::CHAPTERS)?;
|
||||
let edition_pos = ebml::start_master(&mut writer, ebml::EDITION_ENTRY)?;
|
||||
for (i, ch) in chapters.iter().enumerate() {
|
||||
@@ -330,42 +380,60 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
base_pts_ms: None,
|
||||
cues: Vec::new(),
|
||||
frame_count: 0,
|
||||
seek_fixups,
|
||||
info_offset,
|
||||
tracks_offset,
|
||||
chapters_offset,
|
||||
})
|
||||
}
|
||||
|
||||
/// Write a single frame.
|
||||
///
|
||||
/// When `duration_ns` is `Some`, the frame is emitted as a
|
||||
/// `BlockGroup` with `BlockDuration` so the player knows exactly
|
||||
/// when to remove the on-screen artifact (the practical case is
|
||||
/// PGS subtitles — without it, the last bitmap lingers until the
|
||||
/// next display set replaces it). Otherwise a plain `SimpleBlock`.
|
||||
pub fn write_frame(
|
||||
&mut self,
|
||||
track_idx: usize,
|
||||
pts_ns: i64,
|
||||
keyframe: bool,
|
||||
data: &[u8],
|
||||
duration_ns: Option<u64>,
|
||||
) -> io::Result<()> {
|
||||
let raw_ms = pts_ns / 1_000_000;
|
||||
let base = *self.base_pts_ms.get_or_insert(raw_ms);
|
||||
let pts_ms = raw_ms - base;
|
||||
|
||||
// Start new cluster if needed
|
||||
if !self.cluster_open || (pts_ms - self.cluster_ts_ms) >= CLUSTER_DURATION_MS {
|
||||
if self.cluster_open {
|
||||
// Close current cluster (it's a master with unknown size — we use known size)
|
||||
// Actually, for streaming we keep clusters open-ended. Just start a new one.
|
||||
// Cluster boundaries must coincide with a video keyframe so every
|
||||
// Cues entry resolves to a seekable IDR at the cluster start.
|
||||
let is_video_key = keyframe && track_idx == 0;
|
||||
let needs_new_cluster = !self.cluster_open
|
||||
|| (is_video_key && (pts_ms - self.cluster_ts_ms) >= CLUSTER_DURATION_MS);
|
||||
|
||||
if needs_new_cluster {
|
||||
if !is_video_key {
|
||||
return Ok(());
|
||||
}
|
||||
self.start_cluster(pts_ms)?;
|
||||
|
||||
// Add cue point at cluster start for keyframes (video track 0)
|
||||
if keyframe && track_idx == 0 {
|
||||
self.cues.push(CuePoint {
|
||||
timestamp_ms: pts_ms,
|
||||
track: track_idx + 1,
|
||||
cluster_pos: self.cluster_pos - self.segment_start,
|
||||
});
|
||||
}
|
||||
self.cues.push(CuePoint {
|
||||
timestamp_ms: pts_ms,
|
||||
track: track_idx + 1,
|
||||
cluster_pos: self.cluster_pos - self.segment_start,
|
||||
});
|
||||
}
|
||||
|
||||
// Write SimpleBlock
|
||||
let relative_ts = (pts_ms - self.cluster_ts_ms) as i16;
|
||||
self.write_simple_block(track_idx + 1, relative_ts, keyframe, data)?;
|
||||
match duration_ns {
|
||||
Some(dur_ns) => {
|
||||
let duration_ms = (dur_ns / 1_000_000).max(1);
|
||||
self.write_block_group(track_idx + 1, relative_ts, keyframe, data, duration_ms)?;
|
||||
}
|
||||
None => {
|
||||
self.write_simple_block(track_idx + 1, relative_ts, keyframe, data)?;
|
||||
}
|
||||
}
|
||||
self.frame_count += 1;
|
||||
|
||||
Ok(())
|
||||
@@ -377,6 +445,8 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
self.end_cluster()?;
|
||||
|
||||
// Write Cues
|
||||
let cues_start = self.writer.stream_position()?;
|
||||
let cues_offset = cues_start - self.segment_start;
|
||||
if !self.cues.is_empty() {
|
||||
let cues_pos = ebml::start_master(&mut self.writer, ebml::CUES)?;
|
||||
for cue in &self.cues {
|
||||
@@ -395,6 +465,21 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
ebml::end_master(&mut self.writer, cues_pos)?;
|
||||
}
|
||||
|
||||
// Back-patch SeekHead SeekPosition values now that all element offsets are known.
|
||||
for fixup in &self.seek_fixups {
|
||||
let offset = match fixup.target_id {
|
||||
ebml::INFO => self.info_offset,
|
||||
ebml::TRACKS => self.tracks_offset,
|
||||
ebml::CHAPTERS => self.chapters_offset.unwrap_or(0),
|
||||
ebml::CUES => cues_offset,
|
||||
_ => 0,
|
||||
};
|
||||
self.writer
|
||||
.seek(std::io::SeekFrom::Start(fixup.value_offset))?;
|
||||
self.writer.write_all(&offset.to_be_bytes())?;
|
||||
}
|
||||
self.writer.seek(std::io::SeekFrom::End(0))?;
|
||||
|
||||
self.writer.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
@@ -447,6 +532,34 @@ impl<W: Write + Seek> MkvMuxer<W> {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_block_group(
|
||||
&mut self,
|
||||
track_num: usize,
|
||||
relative_ts: i16,
|
||||
keyframe: bool,
|
||||
data: &[u8],
|
||||
duration_ms: u64,
|
||||
) -> io::Result<()> {
|
||||
let track_vint = if track_num < 0x80 {
|
||||
vec![(track_num as u8) | 0x80]
|
||||
} else {
|
||||
vec![0x40 | ((track_num >> 8) as u8), track_num as u8]
|
||||
};
|
||||
let flags: u8 = if keyframe { 0x80 } else { 0x00 };
|
||||
let block_size = track_vint.len() + 2 + 1 + data.len();
|
||||
|
||||
let bg_pos = ebml::start_master(&mut self.writer, ebml::BLOCK_GROUP)?;
|
||||
ebml::write_id(&mut self.writer, ebml::BLOCK)?;
|
||||
ebml::write_size(&mut self.writer, block_size as u64)?;
|
||||
self.writer.write_all(&track_vint)?;
|
||||
self.writer.write_all(&relative_ts.to_be_bytes())?;
|
||||
self.writer.write_all(&[flags])?;
|
||||
self.writer.write_all(data)?;
|
||||
ebml::write_uint(&mut self.writer, ebml::BLOCK_DURATION, duration_ms)?;
|
||||
ebml::end_master(&mut self.writer, bg_pos)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
@@ -556,7 +669,7 @@ mod tests {
|
||||
let tracks = [make_video_track()];
|
||||
let mut muxer = MkvMuxer::new(buf, &tracks, None, 60.0, &[]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 0, true, &[0xDE, 0xAD, 0xBE, 0xEF])
|
||||
.write_frame(0, 0, true, &[0xDE, 0xAD, 0xBE, 0xEF], None)
|
||||
.unwrap();
|
||||
let data = muxer.writer.into_inner();
|
||||
assert!(
|
||||
@@ -594,7 +707,9 @@ mod tests {
|
||||
let writer = SharedWriter(shared.clone());
|
||||
let tracks = [make_video_track()];
|
||||
let mut muxer = MkvMuxer::new(writer, &tracks, Some("Cue Test"), 60.0, &[]).unwrap();
|
||||
muxer.write_frame(0, 0, true, &[0x01, 0x02, 0x03]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 0, true, &[0x01, 0x02, 0x03], None)
|
||||
.unwrap();
|
||||
muxer.finish().unwrap();
|
||||
|
||||
let data = shared.lock().unwrap().clone().into_inner();
|
||||
@@ -610,13 +725,17 @@ mod tests {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let mut muxer = MkvMuxer::new(buf, &tracks, Some("Multi"), 120.0, &[]).unwrap();
|
||||
// Write frames to both tracks
|
||||
muxer.write_frame(0, 0, true, &[0x00, 0x00, 0x01]).unwrap();
|
||||
muxer.write_frame(1, 0, false, &[0x0B, 0x77, 0x00]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 40_000_000, false, &[0x00, 0x00, 0x01])
|
||||
.write_frame(0, 0, true, &[0x00, 0x00, 0x01], None)
|
||||
.unwrap();
|
||||
muxer
|
||||
.write_frame(1, 32_000_000, false, &[0x0B, 0x77, 0x01])
|
||||
.write_frame(1, 0, false, &[0x0B, 0x77, 0x00], None)
|
||||
.unwrap();
|
||||
muxer
|
||||
.write_frame(0, 40_000_000, false, &[0x00, 0x00, 0x01], None)
|
||||
.unwrap();
|
||||
muxer
|
||||
.write_frame(1, 32_000_000, false, &[0x0B, 0x77, 0x01], None)
|
||||
.unwrap();
|
||||
// Should not panic
|
||||
let data = muxer.writer.into_inner();
|
||||
@@ -631,10 +750,12 @@ mod tests {
|
||||
|
||||
// Record position before first frame
|
||||
let pos_before_kf = muxer.writer.position();
|
||||
muxer.write_frame(0, 0, true, &[0xAA]).unwrap();
|
||||
muxer.write_frame(0, 0, true, &[0xAA], None).unwrap();
|
||||
let pos_after_kf = muxer.writer.position();
|
||||
|
||||
muxer.write_frame(0, 1_000_000, false, &[0xBB]).unwrap();
|
||||
muxer
|
||||
.write_frame(0, 1_000_000, false, &[0xBB], None)
|
||||
.unwrap();
|
||||
let pos_after_nkf = muxer.writer.position();
|
||||
|
||||
let data = muxer.writer.into_inner();
|
||||
@@ -812,4 +933,422 @@ mod tests {
|
||||
"FlagForced element should not be present for non-forced subtitle"
|
||||
);
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Seekability tests: SeekHead, keyframe-aligned clusters, Cues
|
||||
// ============================================================
|
||||
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
/// Writer that lets the test inspect the buffer after `finish()` consumes the muxer.
|
||||
struct SharedWriter(Arc<Mutex<Cursor<Vec<u8>>>>);
|
||||
impl Write for SharedWriter {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.0.lock().unwrap().write(buf)
|
||||
}
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
self.0.lock().unwrap().flush()
|
||||
}
|
||||
}
|
||||
impl Seek for SharedWriter {
|
||||
fn seek(&mut self, pos: io::SeekFrom) -> io::Result<u64> {
|
||||
self.0.lock().unwrap().seek(pos)
|
||||
}
|
||||
}
|
||||
|
||||
/// Build interleaved frames at 24 fps video (IDR every gop_secs) + 48 kHz audio (1024 samples per frame).
|
||||
fn frames_for(duration_secs: f64, gop_secs: f64) -> Vec<(usize, i64, bool, Vec<u8>)> {
|
||||
let video_interval_ns: i64 = 1_000_000_000 / 24;
|
||||
let audio_interval_ns: i64 = (1024i64 * 1_000_000_000) / 48_000;
|
||||
let gop_frames = (gop_secs * 24.0).round() as i64;
|
||||
|
||||
let mut out: Vec<(usize, i64, bool, Vec<u8>)> = Vec::new();
|
||||
let total_ns = (duration_secs * 1_000_000_000.0) as i64;
|
||||
|
||||
let mut vi: i64 = 0;
|
||||
loop {
|
||||
let pts = vi * video_interval_ns;
|
||||
if pts >= total_ns {
|
||||
break;
|
||||
}
|
||||
let keyframe = vi % gop_frames == 0;
|
||||
out.push((0, pts, keyframe, vec![0xAB; 64]));
|
||||
vi += 1;
|
||||
}
|
||||
|
||||
let mut ai: i64 = 0;
|
||||
loop {
|
||||
let pts = ai * audio_interval_ns;
|
||||
if pts >= total_ns {
|
||||
break;
|
||||
}
|
||||
out.push((1, pts, true, vec![0xCD; 32]));
|
||||
ai += 1;
|
||||
}
|
||||
|
||||
out.sort_by_key(|f| f.1);
|
||||
out
|
||||
}
|
||||
|
||||
/// Mux frames through a SharedWriter and return the final buffer.
|
||||
fn mux_to_bytes(
|
||||
tracks: &[MkvTrack],
|
||||
chapters: &[Chapter],
|
||||
frames: &[(usize, i64, bool, Vec<u8>)],
|
||||
) -> (Vec<u8>, u64) {
|
||||
let shared = Arc::new(Mutex::new(Cursor::new(Vec::new())));
|
||||
let writer = SharedWriter(shared.clone());
|
||||
let mut muxer = MkvMuxer::new(writer, tracks, None, 0.0, chapters).unwrap();
|
||||
for (t, pts, kf, data) in frames {
|
||||
muxer.write_frame(*t, *pts, *kf, data, None).unwrap();
|
||||
}
|
||||
let frame_count = muxer.frame_count;
|
||||
muxer.finish().unwrap();
|
||||
let data = shared.lock().unwrap().clone().into_inner();
|
||||
(data, frame_count)
|
||||
}
|
||||
|
||||
/// Find the Segment header in the buffer and return (segment_id_pos, segment_start_pos).
|
||||
/// segment_start = position immediately after Segment's id + size bytes.
|
||||
fn locate_segment(data: &[u8]) -> (usize, usize) {
|
||||
let segment_id_pos = find_id(data, ebml::SEGMENT).expect("segment id not found");
|
||||
// Segment is written via write_id + write_unknown_size: 4 byte id + 8 byte size
|
||||
(segment_id_pos, segment_id_pos + 4 + 8)
|
||||
}
|
||||
|
||||
/// Walk Segment's top-level children. Returns Vec<(id, data_start_offset, data_size)>
|
||||
/// where data_start_offset is absolute file offset and data_size is the element body size.
|
||||
fn segment_children(data: &[u8]) -> Vec<(u32, usize, u64)> {
|
||||
let (_, seg_start) = locate_segment(data);
|
||||
let mut out = Vec::new();
|
||||
let mut cursor = Cursor::new(&data[seg_start..]);
|
||||
while (cursor.position() as usize) < data.len() - seg_start {
|
||||
let pos_before = cursor.position();
|
||||
let (id, size, hdr_len) = match ebml::read_element_header(&mut cursor) {
|
||||
Ok(v) => v,
|
||||
Err(_) => break,
|
||||
};
|
||||
let data_abs = seg_start + pos_before as usize + hdr_len;
|
||||
out.push((id, data_abs, size));
|
||||
// Skip the body to advance to the next element.
|
||||
cursor
|
||||
.seek(io::SeekFrom::Current(size as i64))
|
||||
.expect("seek past element body");
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Find every Cluster: returns Vec<(cluster_data_start_abs, cluster_data_size, cluster_timestamp_ms)>.
|
||||
fn find_clusters(data: &[u8]) -> Vec<(usize, u64, u64)> {
|
||||
let mut out = Vec::new();
|
||||
for (id, body_start, body_size) in segment_children(data) {
|
||||
if id == ebml::CLUSTER {
|
||||
let mut cursor = Cursor::new(&data[body_start..body_start + body_size as usize]);
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(
|
||||
tid,
|
||||
ebml::CLUSTER_TIMESTAMP,
|
||||
"cluster must start with timestamp"
|
||||
);
|
||||
let ts = ebml::read_uint_val(&mut cursor, tsize as usize).unwrap();
|
||||
out.push((body_start, body_size, ts));
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Parse the first SimpleBlock that appears in a cluster body slice.
|
||||
/// Returns (track_num, flags_byte). track_num decoded from VINT.
|
||||
fn first_simple_block(cluster_body: &[u8]) -> (u64, u8) {
|
||||
let mut cursor = Cursor::new(cluster_body);
|
||||
loop {
|
||||
let (id, size, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
if id == ebml::SIMPLE_BLOCK {
|
||||
let body_start = cursor.position() as usize;
|
||||
// Decode track VINT.
|
||||
let b0 = cluster_body[body_start];
|
||||
let (track_num, vint_len) = if b0 & 0x80 != 0 {
|
||||
((b0 & 0x7F) as u64, 1usize)
|
||||
} else if b0 & 0x40 != 0 {
|
||||
let b1 = cluster_body[body_start + 1];
|
||||
((((b0 & 0x3F) as u64) << 8) | b1 as u64, 2)
|
||||
} else {
|
||||
panic!("unsupported track vint width");
|
||||
};
|
||||
let flags = cluster_body[body_start + vint_len + 2];
|
||||
return (track_num, flags);
|
||||
}
|
||||
// Skip non-SimpleBlock child.
|
||||
cursor.seek(io::SeekFrom::Current(size as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the Cues element body into Vec<(cue_time, cue_track, cue_cluster_position)>.
|
||||
fn parse_cues(data: &[u8]) -> Vec<(u64, u64, u64)> {
|
||||
let mut out = Vec::new();
|
||||
let (cues_id, cues_body_start, cues_body_size) = segment_children(data)
|
||||
.into_iter()
|
||||
.find(|(id, _, _)| *id == ebml::CUES)
|
||||
.expect("cues element not found");
|
||||
assert_eq!(cues_id, ebml::CUES);
|
||||
let cues_body = &data[cues_body_start..cues_body_start + cues_body_size as usize];
|
||||
let mut cursor = Cursor::new(cues_body);
|
||||
while (cursor.position() as usize) < cues_body.len() {
|
||||
let (id, size, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(id, ebml::CUE_POINT);
|
||||
let cp_end = cursor.position() + size;
|
||||
let mut cue_time = 0u64;
|
||||
let mut cue_track = 0u64;
|
||||
let mut cue_pos = 0u64;
|
||||
while cursor.position() < cp_end {
|
||||
let (sid, ssize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
match sid {
|
||||
ebml::CUE_TIME => {
|
||||
cue_time = ebml::read_uint_val(&mut cursor, ssize as usize).unwrap();
|
||||
}
|
||||
ebml::CUE_TRACK_POSITIONS => {
|
||||
let ctp_end = cursor.position() + ssize;
|
||||
while cursor.position() < ctp_end {
|
||||
let (iid, isize_, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
match iid {
|
||||
ebml::CUE_TRACK => {
|
||||
cue_track =
|
||||
ebml::read_uint_val(&mut cursor, isize_ as usize).unwrap();
|
||||
}
|
||||
ebml::CUE_CLUSTER_POSITION => {
|
||||
cue_pos =
|
||||
ebml::read_uint_val(&mut cursor, isize_ as usize).unwrap();
|
||||
}
|
||||
_ => {
|
||||
cursor.seek(io::SeekFrom::Current(isize_ as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
cursor.seek(io::SeekFrom::Current(ssize as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
out.push((cue_time, cue_track, cue_pos));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Parse the SeekHead body into Vec<(seek_id, seek_position)>.
|
||||
fn parse_seekhead(data: &[u8]) -> Vec<(u32, u64)> {
|
||||
let mut out = Vec::new();
|
||||
let (sh_id, sh_body_start, sh_body_size) = segment_children(data)
|
||||
.into_iter()
|
||||
.find(|(id, _, _)| *id == ebml::SEEK_HEAD)
|
||||
.expect("seekhead not found");
|
||||
assert_eq!(sh_id, ebml::SEEK_HEAD);
|
||||
let sh_body = &data[sh_body_start..sh_body_start + sh_body_size as usize];
|
||||
let mut cursor = Cursor::new(sh_body);
|
||||
while (cursor.position() as usize) < sh_body.len() {
|
||||
let (id, size, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(id, ebml::SEEK);
|
||||
let seek_end = cursor.position() + size;
|
||||
let mut seek_id_val: u32 = 0;
|
||||
let mut seek_pos_val: u64 = 0;
|
||||
while cursor.position() < seek_end {
|
||||
let (sid, ssize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
match sid {
|
||||
ebml::SEEK_ID => {
|
||||
let raw = ebml::read_uint_val(&mut cursor, ssize as usize).unwrap();
|
||||
seek_id_val = raw as u32;
|
||||
}
|
||||
ebml::SEEK_POSITION => {
|
||||
seek_pos_val = ebml::read_uint_val(&mut cursor, ssize as usize).unwrap();
|
||||
}
|
||||
_ => {
|
||||
cursor.seek(io::SeekFrom::Current(ssize as i64)).unwrap();
|
||||
}
|
||||
}
|
||||
}
|
||||
out.push((seek_id_val, seek_pos_val));
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cluster_starts_only_on_video_keyframe() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let clusters = find_clusters(&data);
|
||||
assert!(!clusters.is_empty(), "expected at least one cluster");
|
||||
for (body_start, body_size, _ts) in clusters {
|
||||
let body = &data[body_start..body_start + body_size as usize];
|
||||
// Skip past the CLUSTER_TIMESTAMP element first.
|
||||
let mut cursor = Cursor::new(body);
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap();
|
||||
let after_ts = cursor.position() as usize;
|
||||
let (track_num, flags) = first_simple_block(&body[after_ts..]);
|
||||
assert_eq!(
|
||||
track_num, 1,
|
||||
"first block in cluster must be track 1 (video)"
|
||||
);
|
||||
assert_eq!(
|
||||
flags & 0x80,
|
||||
0x80,
|
||||
"first block in cluster must have keyframe flag set, got 0x{:02X}",
|
||||
flags
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cue_count_equals_cluster_count() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let clusters = find_clusters(&data);
|
||||
let cues = parse_cues(&data);
|
||||
assert_eq!(
|
||||
clusters.len(),
|
||||
cues.len(),
|
||||
"cluster count {} != cue count {}",
|
||||
clusters.len(),
|
||||
cues.len()
|
||||
);
|
||||
// For 30s @ 5s min cluster duration with 1s GOP, expect 6 clusters / 6 cues.
|
||||
assert_eq!(
|
||||
clusters.len(),
|
||||
6,
|
||||
"expected 6 clusters for 30s @ 5s cluster duration"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cue_positions_resolve_to_clusters() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let (_, seg_start) = locate_segment(&data);
|
||||
let cues = parse_cues(&data);
|
||||
assert!(!cues.is_empty());
|
||||
for (_time, _track, pos) in cues {
|
||||
let abs = seg_start + pos as usize;
|
||||
let mut cursor = Cursor::new(&data[abs..]);
|
||||
let (id, _size, _hdr_len) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(
|
||||
id,
|
||||
ebml::CLUSTER,
|
||||
"cue position 0x{:X} did not resolve to a cluster",
|
||||
pos
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cue_times_match_cluster_timestamps() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let frames = frames_for(30.0, 1.0);
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames);
|
||||
let (_, seg_start) = locate_segment(&data);
|
||||
let cues = parse_cues(&data);
|
||||
for (time, _track, pos) in cues {
|
||||
let abs = seg_start + pos as usize;
|
||||
let mut cursor = Cursor::new(&data[abs..]);
|
||||
let (id, size, _hdr_len) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(id, ebml::CLUSTER);
|
||||
let body_start = abs + (cursor.position() as usize);
|
||||
let body = &data[body_start..body_start + size as usize];
|
||||
let mut bc = Cursor::new(body);
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut bc).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
let cluster_ts = ebml::read_uint_val(&mut bc, tsize as usize).unwrap();
|
||||
assert_eq!(
|
||||
cluster_ts, time,
|
||||
"cluster timestamp {} != cue time {}",
|
||||
cluster_ts, time
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seekhead_is_first_child_of_segment() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames_for(10.0, 1.0));
|
||||
let children = segment_children(&data);
|
||||
assert!(!children.is_empty());
|
||||
assert_eq!(
|
||||
children[0].0,
|
||||
ebml::SEEK_HEAD,
|
||||
"first child of segment must be SeekHead, got id 0x{:X}",
|
||||
children[0].0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seekhead_points_to_real_elements() {
|
||||
let tracks = [make_video_track(), make_audio_track()];
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames_for(10.0, 1.0));
|
||||
let (_, seg_start) = locate_segment(&data);
|
||||
let entries = parse_seekhead(&data);
|
||||
let required = [ebml::INFO, ebml::TRACKS, ebml::CUES];
|
||||
for &want_id in &required {
|
||||
let entry = entries
|
||||
.iter()
|
||||
.find(|(id, _)| *id == want_id)
|
||||
.unwrap_or_else(|| panic!("seekhead missing entry for id 0x{:X}", want_id));
|
||||
let abs = seg_start + entry.1 as usize;
|
||||
let mut cursor = Cursor::new(&data[abs..]);
|
||||
let (got_id, _, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(
|
||||
got_id, want_id,
|
||||
"seekhead entry for 0x{:X} resolves to wrong id 0x{:X}",
|
||||
want_id, got_id
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn seekhead_omits_chapters_when_empty() {
|
||||
let tracks = [make_video_track()];
|
||||
let (data, _) = mux_to_bytes(&tracks, &[], &frames_for(5.0, 1.0));
|
||||
let entries = parse_seekhead(&data);
|
||||
assert_eq!(
|
||||
entries.len(),
|
||||
3,
|
||||
"expected 3 seek entries (Info, Tracks, Cues), got {}",
|
||||
entries.len()
|
||||
);
|
||||
assert!(
|
||||
entries.iter().all(|(id, _)| *id != ebml::CHAPTERS),
|
||||
"seekhead should not contain Chapters entry when chapters are empty"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pre_first_keyframe_frames_dropped() {
|
||||
let tracks = [make_video_track()];
|
||||
let frames = vec![
|
||||
(0usize, 0i64, false, vec![0x11; 16]),
|
||||
(0usize, 41_000_000i64, true, vec![0x22; 16]),
|
||||
];
|
||||
let (data, frame_count) = mux_to_bytes(&tracks, &[], &frames);
|
||||
assert_eq!(frame_count, 1, "muxer.frame_count must equal 1");
|
||||
let clusters = find_clusters(&data);
|
||||
assert_eq!(clusters.len(), 1, "expected exactly one cluster");
|
||||
let (body_start, body_size, _ts) = clusters[0];
|
||||
let body = &data[body_start..body_start + body_size as usize];
|
||||
let mut cursor = Cursor::new(body);
|
||||
// Skip CLUSTER_TIMESTAMP.
|
||||
let (tid, tsize, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
assert_eq!(tid, ebml::CLUSTER_TIMESTAMP);
|
||||
cursor.seek(io::SeekFrom::Current(tsize as i64)).unwrap();
|
||||
let mut sb_count = 0;
|
||||
while (cursor.position() as usize) < body.len() {
|
||||
let (id, sz, _) = ebml::read_element_header(&mut cursor).unwrap();
|
||||
if id == ebml::SIMPLE_BLOCK {
|
||||
sb_count += 1;
|
||||
}
|
||||
cursor.seek(io::SeekFrom::Current(sz as i64)).unwrap();
|
||||
}
|
||||
assert_eq!(sb_count, 1, "expected exactly one SimpleBlock in output");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -138,6 +138,7 @@ impl crate::pes::Stream for MkvStream {
|
||||
pts: pts_ms * 1_000_000, // ms → ns
|
||||
keyframe,
|
||||
data,
|
||||
duration_ns: None,
|
||||
}));
|
||||
}
|
||||
_ => {
|
||||
@@ -150,9 +151,13 @@ impl crate::pes::Stream for MkvStream {
|
||||
|
||||
fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> {
|
||||
match &mut self.mode {
|
||||
Mode::Write { muxer: Some(m) } => {
|
||||
m.write_frame(frame.track, frame.pts, frame.keyframe, &frame.data)
|
||||
}
|
||||
Mode::Write { muxer: Some(m) } => m.write_frame(
|
||||
frame.track,
|
||||
frame.pts,
|
||||
frame.keyframe,
|
||||
&frame.data,
|
||||
frame.duration_ns,
|
||||
),
|
||||
Mode::Write { muxer: None } => Ok(()),
|
||||
Mode::Read(_) => Err(crate::error::Error::StreamReadOnly.into()),
|
||||
}
|
||||
|
||||
+15
-2
@@ -17,8 +17,9 @@
|
||||
|
||||
// Public modules — types here are intentionally part of the consumable API.
|
||||
pub mod codec;
|
||||
pub mod demux_thread;
|
||||
pub mod disc;
|
||||
pub mod iso;
|
||||
pub mod pipelined_stream;
|
||||
pub mod resolve;
|
||||
|
||||
// Internal modules — implementation details. Their *types* are re-exported
|
||||
@@ -33,6 +34,17 @@ pub(crate) mod m2ts;
|
||||
/// to round-trip codec_privates that don't fit inside the underlying format).
|
||||
/// Exposed for integration tests that exercise the wire format directly.
|
||||
pub mod meta;
|
||||
|
||||
// ── Phase 3 sequential muxers ──────────────────────────────────────────────
|
||||
//
|
||||
// New container muxers that consume PES frames and write to a
|
||||
// `SequentialSink`. They are NOT refactors of the existing `MkvStream` /
|
||||
// `M2tsStream` (which round-trip via the legacy `Stream` trait + the
|
||||
// BD-TS framing); they're sequential-only and target the Phase 2 sink
|
||||
// split end-to-end.
|
||||
pub mod fmp4;
|
||||
pub mod hevc;
|
||||
pub mod m2ts_mux;
|
||||
pub(crate) mod mkv;
|
||||
pub(crate) mod mkvstream;
|
||||
pub(crate) mod network;
|
||||
@@ -43,11 +55,12 @@ pub(crate) mod ts;
|
||||
pub(crate) mod tsmux;
|
||||
|
||||
pub use disc::DiscStream;
|
||||
pub use iso::IsoSectorReader;
|
||||
pub use m2ts::M2tsStream;
|
||||
pub use mkvstream::MkvStream;
|
||||
pub use network::NetworkStream;
|
||||
pub use null::NullStream;
|
||||
pub use pipelined_stream::PipelinedPesStream;
|
||||
pub use resolve::build_iso_pipeline;
|
||||
pub use resolve::{InputOptions, StreamUrl, input, output, parse_url};
|
||||
pub use stdio::StdioStream;
|
||||
|
||||
|
||||
@@ -186,6 +186,7 @@ mod tests {
|
||||
pts: 90000,
|
||||
keyframe: true,
|
||||
data: vec![0x47; 192],
|
||||
duration_ns: None,
|
||||
};
|
||||
pes::Stream::write(&mut writer, &frame).unwrap();
|
||||
pes::Stream::finish(&mut writer).unwrap();
|
||||
|
||||
@@ -47,6 +47,7 @@ mod tests {
|
||||
pts: 0,
|
||||
keyframe: true,
|
||||
data: vec![0x01, 0x02, 0x03],
|
||||
duration_ns: None,
|
||||
};
|
||||
sink.write(&frame).unwrap();
|
||||
let _ = sink.info();
|
||||
|
||||
@@ -0,0 +1,226 @@
|
||||
//! `PipelinedPesStream` — the read-side of the freemkv mux
|
||||
//! highway.
|
||||
//!
|
||||
//! Given a [`crate::mux::demux_thread::DemuxThread`] (which has the
|
||||
//! producer + demux workers already spawned), a set of codec
|
||||
//! parsers, and the title metadata, this struct implements
|
||||
//! [`crate::pes::Stream`] by running codec parse on the caller's
|
||||
//! thread and emitting `PesFrame`s one at a time.
|
||||
//!
|
||||
//! The pipeline runs three threads in parallel:
|
||||
//!
|
||||
//! ```text
|
||||
//! Thread A: read + decrypt (PrefetchedSectorSource / BytePrefetcher)
|
||||
//! Thread B: M2TS demux (DemuxThread)
|
||||
//! Thread C: codec parse (this struct, on the caller's thread)
|
||||
//! ```
|
||||
//!
|
||||
//! Communication between A→B and B→C is via bounded channels with
|
||||
//! recycled buffer pools — no allocations or memcpys in the steady-
|
||||
//! state hot loop.
|
||||
//!
|
||||
//! This is the *only* read-side `Stream` impl in tree. Both ISO file
|
||||
//! mux ([`crate::mux::resolve`]) and BD-TS file mux ([`crate::mux::M2tsStream`])
|
||||
//! return a `PipelinedPesStream`; the differences are in how the
|
||||
//! producer thread (A) is configured — sector-aligned reads with
|
||||
//! AACS decrypt for ISO, raw byte reads for M2TS.
|
||||
|
||||
use super::codec::CodecParser;
|
||||
use super::demux_thread::{DemuxBatch, DemuxThread};
|
||||
use super::ts::PesPacket;
|
||||
use crate::disc::DiscTitle;
|
||||
use crate::pes::{PesFrame, Stream};
|
||||
use crossbeam_channel::Receiver;
|
||||
use std::io;
|
||||
|
||||
/// Stream impl that consumes pre-demuxed `PesPacket` batches from a
|
||||
/// [`DemuxThread`] and runs codec parse on the caller's thread.
|
||||
pub struct PipelinedPesStream {
|
||||
title: DiscTitle,
|
||||
parsers: Vec<(u16, Box<dyn CodecParser>)>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
|
||||
demux_rx: Receiver<DemuxBatch>,
|
||||
/// Kept alive so dropping this stream joins the demux + producer
|
||||
/// workers deterministically. Never poked directly after spawn.
|
||||
#[allow(dead_code)]
|
||||
demux_thread: DemuxThread,
|
||||
|
||||
pending_frames: std::collections::VecDeque<PesFrame>,
|
||||
eof: bool,
|
||||
}
|
||||
|
||||
impl PipelinedPesStream {
|
||||
/// Wire up the stream. Caller has already spawned the
|
||||
/// `DemuxThread` (which in turn owns the producer); we take the
|
||||
/// receiver end + the join handle bundle so cleanup is bounded
|
||||
/// on drop.
|
||||
pub fn new(
|
||||
demux_thread: DemuxThread,
|
||||
demux_rx: Receiver<DemuxBatch>,
|
||||
title: DiscTitle,
|
||||
parsers: Vec<(u16, Box<dyn CodecParser>)>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
) -> Self {
|
||||
Self {
|
||||
title,
|
||||
parsers,
|
||||
pid_to_track,
|
||||
demux_rx,
|
||||
demux_thread,
|
||||
pending_frames: std::collections::VecDeque::new(),
|
||||
eof: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Pull one batch of `PesPacket`s from the demux thread, run
|
||||
/// codec parse on each, enqueue resulting `PesFrame`s on
|
||||
/// `pending_frames`. Returns Ok(true) on success, Ok(false) on
|
||||
/// EOF (channel closed cleanly), Err on demuxer error.
|
||||
fn pump_one_batch(&mut self) -> io::Result<bool> {
|
||||
match self.demux_rx.recv() {
|
||||
Ok(DemuxBatch::Ts(packets)) => {
|
||||
self.consume_ts(packets);
|
||||
Ok(true)
|
||||
}
|
||||
Ok(DemuxBatch::Ps(packets)) => {
|
||||
self.consume_ps(packets);
|
||||
Ok(true)
|
||||
}
|
||||
Ok(DemuxBatch::Err(e)) => Err(e),
|
||||
Err(_) => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
fn consume_ts(&mut self, packets: Vec<PesPacket>) {
|
||||
let skip_parse = std::env::var_os("FREEMKV_SKIP_PARSE").is_some();
|
||||
for pes in packets {
|
||||
if let Some((_, track)) = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(pid, _)| *pid == pes.pid)
|
||||
.copied()
|
||||
{
|
||||
if skip_parse {
|
||||
// Profiling escape hatch — bypass codec parser.
|
||||
self.pending_frames.push_back(PesFrame {
|
||||
track,
|
||||
pts: pes.pts.map(super::codec::pts_to_ns).unwrap_or(0),
|
||||
keyframe: false,
|
||||
data: pes.data,
|
||||
duration_ns: None,
|
||||
});
|
||||
} else if let Some((_, parser)) =
|
||||
self.parsers.iter_mut().find(|(pid, _)| *pid == pes.pid)
|
||||
{
|
||||
for frame in parser.parse(&pes) {
|
||||
self.pending_frames
|
||||
.push_back(PesFrame::from_codec_frame(track, frame));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn consume_ps(&mut self, packets: Vec<super::ps::PsPacket>) {
|
||||
for ps in packets {
|
||||
let track = match ps.stream_id {
|
||||
0xE0..=0xEF => 0,
|
||||
0xC0..=0xDF => 1,
|
||||
0xBD => ps
|
||||
.sub_stream_id
|
||||
.map(|s| (s & 0x1F) as usize + 1)
|
||||
.unwrap_or(1),
|
||||
_ => continue,
|
||||
};
|
||||
if track >= self.title.streams.len() {
|
||||
continue;
|
||||
}
|
||||
let pid = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(_, idx)| *idx == track)
|
||||
.map(|(p, _)| *p)
|
||||
.unwrap_or(0);
|
||||
let pes = PesPacket {
|
||||
pid,
|
||||
pts: ps.pts.map(|p| p as i64),
|
||||
dts: ps.dts.map(|d| d as i64),
|
||||
data: ps.data,
|
||||
};
|
||||
if let Some((_, parser)) = self.parsers.iter_mut().find(|(p, _)| *p == pid) {
|
||||
for frame in parser.parse(&pes) {
|
||||
self.pending_frames
|
||||
.push_back(PesFrame::from_codec_frame(track, frame));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Stream for PipelinedPesStream {
|
||||
fn read(&mut self) -> io::Result<Option<PesFrame>> {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
if self.eof {
|
||||
return Ok(None);
|
||||
}
|
||||
loop {
|
||||
match self.pump_one_batch()? {
|
||||
true => {
|
||||
if let Some(frame) = self.pending_frames.pop_front() {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
// Batch contained no trackable packets — pull again.
|
||||
}
|
||||
false => {
|
||||
self.eof = true;
|
||||
return Ok(self.pending_frames.pop_front());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn write(&mut self, _: &PesFrame) -> io::Result<()> {
|
||||
Err(crate::error::Error::StreamReadOnly.into())
|
||||
}
|
||||
|
||||
fn finish(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn info(&self) -> &DiscTitle {
|
||||
&self.title
|
||||
}
|
||||
|
||||
fn headers_ready(&self) -> bool {
|
||||
// Match the previous DiscStream semantics: video tracks need
|
||||
// codec_private before the consumer can write the container
|
||||
// header. FREEMKV_SKIP_PARSE forces ready (no parser ever
|
||||
// populates codec_private in that mode).
|
||||
if std::env::var_os("FREEMKV_SKIP_PARSE").is_some() {
|
||||
return true;
|
||||
}
|
||||
for (idx, s) in self.title.streams.iter().enumerate() {
|
||||
if let crate::disc::Stream::Video(v) = s {
|
||||
if !v.secondary && self.codec_private(idx).is_none() {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
|
||||
let pid = self
|
||||
.pid_to_track
|
||||
.iter()
|
||||
.find(|(_, idx)| *idx == track)
|
||||
.map(|(p, _)| *p)?;
|
||||
self.parsers
|
||||
.iter()
|
||||
.find(|(p, _)| *p == pid)
|
||||
.and_then(|(_, parser)| parser.codec_private())
|
||||
}
|
||||
}
|
||||
+170
-10
@@ -15,11 +15,13 @@
|
||||
//! Bare paths without a scheme are rejected.
|
||||
//! For disc→ISO (raw sector copy), use `Disc::copy()` instead.
|
||||
|
||||
use super::disc::DiscStream;
|
||||
use super::network::NetworkStream;
|
||||
use super::null::NullStream;
|
||||
use super::pipelined_stream::PipelinedPesStream;
|
||||
use super::stdio::StdioStream;
|
||||
use super::{M2tsStream, MkvStream};
|
||||
use crate::disc::{ContentFormat, DiscTitle};
|
||||
use crate::sector::SectorSource;
|
||||
use std::io;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
@@ -185,7 +187,12 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::S
|
||||
},
|
||||
None => crate::disc::ScanOptions::default(),
|
||||
};
|
||||
let mut reader = super::iso::IsoSectorReader::open(&path.to_string_lossy())?;
|
||||
// FileSectorSource is the sole file-backed sector source.
|
||||
// It carries the platform-tuned SEQUENTIAL fadvise hint
|
||||
// (so the kernel readahead window widens) and the periodic
|
||||
// DONTNEED page-cache eviction that bounds memory pressure
|
||||
// when the mux output is being written to the same disk.
|
||||
let mut reader = crate::io::file_sector_source::FileSectorSource::open(path)?;
|
||||
let capacity = reader.capacity_sectors();
|
||||
let disc = crate::disc::Disc::scan_image(&mut reader, capacity, &scan_opts)
|
||||
.map_err(|e| -> io::Error { e.into() })?;
|
||||
@@ -203,21 +210,40 @@ pub fn input(url: &str, opts: &InputOptions) -> io::Result<Box<dyn crate::pes::S
|
||||
let title = disc.titles[idx].clone();
|
||||
let keys = disc.decrypt_keys();
|
||||
let format = disc.content_format;
|
||||
// ISO file: use large batch size (16 MB) — sequential read from fast storage, no bad sectors.
|
||||
// Physical drives need small batches for adaptive error handling and retry logic.
|
||||
// ISO file: 16 MiB batch — sequential read from fast
|
||||
// storage, no bad sectors. Measured optimum on the rip1
|
||||
// testbed; bumping to 32 MiB regressed (more cache
|
||||
// pressure, longer per-batch latency starves the consumer
|
||||
// between iterations). Physical drives keep smaller
|
||||
// batches for adaptive error handling.
|
||||
const ISO_MUX_BATCH_SECTORS: u16 = 8192;
|
||||
let mut stream =
|
||||
DiscStream::new(Box::new(reader), title, keys, ISO_MUX_BATCH_SECTORS, format);
|
||||
if opts.raw {
|
||||
stream.set_raw();
|
||||
}
|
||||
|
||||
// Pass `DecryptKeys::None` to the decrypt decorator when
|
||||
// --raw is set — the read stack still flows through the
|
||||
// same producer+demux+parse pipeline, just without the
|
||||
// AACS / CSS step. Single highway for all ISO reads.
|
||||
let effective_keys = if opts.raw {
|
||||
crate::decrypt::DecryptKeys::None
|
||||
} else {
|
||||
keys
|
||||
};
|
||||
let stream = build_iso_pipeline(
|
||||
reader,
|
||||
title,
|
||||
effective_keys,
|
||||
ISO_MUX_BATCH_SECTORS,
|
||||
format,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
Ok(Box::new(stream))
|
||||
}
|
||||
StreamUrl::M2ts { ref path } => {
|
||||
validate_file_path(path, "m2ts")?;
|
||||
let file = std::fs::File::open(path)?;
|
||||
let reader = std::io::BufReader::with_capacity(IO_BUF_SIZE, file);
|
||||
Ok(Box::new(M2tsStream::open(reader)?))
|
||||
let stream = build_m2ts_pipeline(reader)?;
|
||||
Ok(Box::new(stream))
|
||||
}
|
||||
StreamUrl::Mkv { ref path } => {
|
||||
validate_file_path(path, "mkv")?;
|
||||
@@ -281,3 +307,137 @@ pub fn output(
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Demuxer-side state derived from a `DiscTitle`: the codec parser
|
||||
/// table (keyed by PID), the PID-to-track index map, and an initial
|
||||
/// `TsDemuxer` / `PsDemuxer` (whichever the content format calls
|
||||
/// for).
|
||||
type DemuxState = (
|
||||
Vec<(u16, Box<dyn super::codec::CodecParser>)>,
|
||||
Vec<(u16, usize)>,
|
||||
Option<super::ts::TsDemuxer>,
|
||||
Option<super::ps::PsDemuxer>,
|
||||
);
|
||||
|
||||
/// Build the title's codec parser table + initial `TsDemuxer` /
|
||||
/// `PsDemuxer`. Used by both the ISO and M2TS pipeline builders.
|
||||
fn build_demux_state(title: &DiscTitle, format: ContentFormat) -> DemuxState {
|
||||
let mut pids = Vec::new();
|
||||
let mut parsers = Vec::new();
|
||||
let mut pid_to_track = Vec::new();
|
||||
for (idx, s) in title.streams.iter().enumerate() {
|
||||
let (pid, codec) = match s {
|
||||
crate::disc::Stream::Video(v) => (v.pid, v.codec),
|
||||
crate::disc::Stream::Audio(a) => (a.pid, a.codec),
|
||||
crate::disc::Stream::Subtitle(s) => (s.pid, s.codec),
|
||||
};
|
||||
pids.push(pid);
|
||||
pid_to_track.push((pid, idx));
|
||||
parsers.push((pid, super::codec::parser_for_codec(codec, None)));
|
||||
}
|
||||
let (ts, ps) = match format {
|
||||
ContentFormat::MpegPs => (None, Some(super::ps::PsDemuxer::new())),
|
||||
ContentFormat::BdTs => {
|
||||
if pids.is_empty() {
|
||||
(None, None)
|
||||
} else {
|
||||
(Some(super::ts::TsDemuxer::new(&pids)), None)
|
||||
}
|
||||
}
|
||||
};
|
||||
(parsers, pid_to_track, ts, ps)
|
||||
}
|
||||
|
||||
/// Assemble the ISO mux pipeline (read+decrypt → demux → parse) for
|
||||
/// a `FileSectorSource`-backed reader. Returns the resulting
|
||||
/// `PipelinedPesStream`.
|
||||
pub fn build_iso_pipeline<S: SectorSource + Send + 'static>(
|
||||
reader: S,
|
||||
title: DiscTitle,
|
||||
keys: crate::decrypt::DecryptKeys,
|
||||
batch_sectors: u16,
|
||||
format: ContentFormat,
|
||||
halt: Option<crate::halt::Halt>,
|
||||
event_fn: Option<crate::sector::prefetched::EventFn>,
|
||||
) -> PipelinedPesStream {
|
||||
let extents = title.extents.clone();
|
||||
let decrypting =
|
||||
crate::sector::DecryptingSectorSource::new(Box::new(reader) as Box<dyn SectorSource>, keys);
|
||||
let prefetched = crate::sector::PrefetchedSectorSource::new_with_events(
|
||||
decrypting,
|
||||
extents,
|
||||
batch_sectors,
|
||||
halt.clone(),
|
||||
event_fn,
|
||||
);
|
||||
let (rx, recycle_tx, shell) = prefetched.into_channels();
|
||||
|
||||
let (parsers, pid_to_track, ts, ps) = build_demux_state(&title, format);
|
||||
let (demux_thread, demux_rx) =
|
||||
super::demux_thread::DemuxThread::spawn_zero_copy(rx, recycle_tx, shell, halt, ts, ps);
|
||||
PipelinedPesStream::new(demux_thread, demux_rx, title, parsers, pid_to_track)
|
||||
}
|
||||
|
||||
/// Assemble the M2TS file mux pipeline (read → demux → parse) for a
|
||||
/// byte-stream reader. Scans the head for FMKV header or PMT/PAT,
|
||||
/// rebuilds the title metadata, then wraps a chained reader (head +
|
||||
/// remainder) in a `BytePrefetcher` feeding the demux + parse
|
||||
/// threads.
|
||||
fn build_m2ts_pipeline<R: std::io::Read + Send + 'static>(
|
||||
mut reader: R,
|
||||
) -> io::Result<PipelinedPesStream> {
|
||||
use super::meta;
|
||||
use std::io::Read;
|
||||
|
||||
const M2TS_SCAN_BYTES: usize = 1024 * 1024;
|
||||
let mut head = vec![0u8; M2TS_SCAN_BYTES];
|
||||
let head_len = {
|
||||
let mut filled = 0;
|
||||
while filled < head.len() {
|
||||
match reader.read(&mut head[filled..])? {
|
||||
0 => break,
|
||||
n => filled += n,
|
||||
}
|
||||
}
|
||||
filled
|
||||
};
|
||||
head.truncate(head_len);
|
||||
|
||||
// Try FMKV metadata header first; fall back to PMT scan.
|
||||
let mut cursor = io::Cursor::new(&head);
|
||||
let (title, head_consumed) = if let Ok(Some(m)) = meta::read_header(&mut cursor) {
|
||||
(m.to_title(), cursor.position() as usize)
|
||||
} else {
|
||||
let streams = super::ts::scan_streams(&head)
|
||||
.ok_or_else(|| -> io::Error { crate::error::Error::NoStreams.into() })?;
|
||||
let t = DiscTitle {
|
||||
duration_secs: 0.0,
|
||||
streams,
|
||||
..DiscTitle::empty()
|
||||
};
|
||||
(t, 0)
|
||||
};
|
||||
|
||||
// Chain: any un-consumed head bytes + the remainder of the
|
||||
// reader. The demuxer sees a contiguous M2TS byte stream.
|
||||
let remaining_head = head[head_consumed..].to_vec();
|
||||
let chained: Box<dyn Read + Send> = Box::new(io::Cursor::new(remaining_head).chain(reader));
|
||||
|
||||
let prefetcher = crate::io::byte_prefetcher::BytePrefetcher::new(
|
||||
chained,
|
||||
crate::io::byte_prefetcher::DEFAULT_CHUNK_BYTES,
|
||||
None,
|
||||
);
|
||||
let (rx, recycle_tx, shell) = prefetcher.into_channels();
|
||||
|
||||
let (parsers, pid_to_track, ts, ps) = build_demux_state(&title, ContentFormat::BdTs);
|
||||
let (demux_thread, demux_rx) =
|
||||
super::demux_thread::DemuxThread::spawn_zero_copy(rx, recycle_tx, shell, None, ts, ps);
|
||||
Ok(PipelinedPesStream::new(
|
||||
demux_thread,
|
||||
demux_rx,
|
||||
title,
|
||||
parsers,
|
||||
pid_to_track,
|
||||
))
|
||||
}
|
||||
|
||||
+108
-83
@@ -37,11 +37,20 @@ struct PesAssembler {
|
||||
active: bool,
|
||||
}
|
||||
|
||||
/// Initial capacity for a fresh PES buffer. Sized to cover the
|
||||
/// common BD-TS audio / subtitle PES outright (a few KB to ~16 KB).
|
||||
/// Video PES (typically 150–300 KB on UHD) will grow this via the
|
||||
/// standard Vec doubling, but the doublings hit the allocator's
|
||||
/// slab caches instead of the 64-page first-touch faults that the
|
||||
/// previous `Vec::with_capacity(256 * 1024)` triggered on every PES
|
||||
/// boundary.
|
||||
const PES_BUFFER_INIT_CAP: usize = 16 * 1024;
|
||||
|
||||
impl PesAssembler {
|
||||
fn new(pid: u16) -> Self {
|
||||
Self {
|
||||
pid,
|
||||
buffer: Vec::with_capacity(256 * 1024),
|
||||
buffer: Vec::with_capacity(PES_BUFFER_INIT_CAP),
|
||||
pts: None,
|
||||
dts: None,
|
||||
active: false,
|
||||
@@ -55,7 +64,7 @@ impl PesAssembler {
|
||||
pid: self.pid,
|
||||
pts: self.pts,
|
||||
dts: self.dts,
|
||||
data: std::mem::replace(&mut self.buffer, Vec::with_capacity(256 * 1024)),
|
||||
data: std::mem::replace(&mut self.buffer, Vec::with_capacity(PES_BUFFER_INIT_CAP)),
|
||||
})
|
||||
} else {
|
||||
self.buffer.clear();
|
||||
@@ -124,97 +133,54 @@ impl TsDemuxer {
|
||||
}
|
||||
}
|
||||
|
||||
/// Feed a chunk of BD transport stream data. Handles non-192-byte-aligned input
|
||||
/// by buffering leftover bytes between calls. Returns completed PES packets.
|
||||
/// Feed a chunk of BD transport stream data. Handles non-192-byte-
|
||||
/// aligned input by buffering leftover bytes between calls. Returns
|
||||
/// completed PES packets.
|
||||
///
|
||||
/// 16 MiB ISO batches never divide evenly into 192-byte BD-TS
|
||||
/// packets, so every call after the first carries a ~64-byte
|
||||
/// remainder. The pre-0.24 implementation handled this by building
|
||||
/// a `combined` Vec containing remainder + the entire new input —
|
||||
/// a 16 MiB+ memcpy on every call. Now we splice exactly one
|
||||
/// boundary packet from a stack buffer, then process the rest of
|
||||
/// `data` in place. Zero-copy on the bulk path; one 192-byte copy
|
||||
/// on the boundary.
|
||||
pub fn feed(&mut self, data: &[u8]) -> Vec<PesPacket> {
|
||||
let mut completed = Vec::with_capacity(4);
|
||||
|
||||
// Prepend any remainder from previous call
|
||||
let mut combined: Vec<u8> = Vec::new();
|
||||
let work: &[u8] = if !self.remainder.is_empty() {
|
||||
combined.reserve(self.remainder.len() + data.len());
|
||||
combined.extend_from_slice(&self.remainder);
|
||||
combined.extend_from_slice(data);
|
||||
self.remainder.clear();
|
||||
&combined
|
||||
} else {
|
||||
data
|
||||
};
|
||||
|
||||
let mut offset = 0;
|
||||
|
||||
while offset + BD_TS_PACKET_SIZE <= work.len() {
|
||||
let packet = &work[offset..offset + BD_TS_PACKET_SIZE];
|
||||
offset += BD_TS_PACKET_SIZE;
|
||||
|
||||
// Skip 4-byte TP_extra_header, check sync byte
|
||||
if packet[4] != SYNC_BYTE {
|
||||
continue;
|
||||
}
|
||||
|
||||
let ts = &packet[4..]; // 188-byte standard TS packet
|
||||
|
||||
// Parse TS header
|
||||
let pid = (((ts[1] & 0x1F) as u16) << 8) | ts[2] as u16;
|
||||
let pusi = ts[1] & 0x40 != 0; // Payload Unit Start Indicator
|
||||
let adaptation = (ts[3] >> 4) & 0x03;
|
||||
|
||||
// Check if we're tracking this PID
|
||||
let idx = if (pid as usize) < self.pid_index.len() {
|
||||
self.pid_index[pid as usize]
|
||||
} else {
|
||||
-1
|
||||
};
|
||||
if idx < 0 {
|
||||
continue;
|
||||
}
|
||||
let asm = &mut self.assemblers[idx as usize];
|
||||
|
||||
// Find payload start (skip adaptation field if present)
|
||||
let payload_start = if adaptation == 0x03 || adaptation == 0x02 {
|
||||
// Adaptation field present
|
||||
let af_len = ts[4] as usize;
|
||||
if af_len > 183 {
|
||||
continue; // Malformed: AF length exceeds TS payload
|
||||
}
|
||||
5 + af_len
|
||||
} else {
|
||||
4
|
||||
};
|
||||
|
||||
if payload_start >= TS_PACKET_SIZE {
|
||||
continue;
|
||||
}
|
||||
|
||||
// No payload
|
||||
if adaptation == 0x02 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let payload = &ts[payload_start..];
|
||||
|
||||
if pusi {
|
||||
// New PES packet starts here — parse PES header
|
||||
let (pts, dts, pes_data_start) = parse_pes_header(payload);
|
||||
if let Some(prev) = asm.start(pts, dts) {
|
||||
completed.push(prev);
|
||||
}
|
||||
if pes_data_start < payload.len() {
|
||||
asm.push(&payload[pes_data_start..]);
|
||||
}
|
||||
} else {
|
||||
// Continuation of current PES packet
|
||||
asm.push(payload);
|
||||
// Boundary packet: if a partial packet was left from the last
|
||||
// call, complete it from the head of `data` without touching
|
||||
// the rest of `data`.
|
||||
if !self.remainder.is_empty() {
|
||||
let need = BD_TS_PACKET_SIZE - self.remainder.len();
|
||||
if data.len() < need {
|
||||
// Still not a full packet — accumulate and wait.
|
||||
self.remainder.extend_from_slice(data);
|
||||
return completed;
|
||||
}
|
||||
let mut boundary = [0u8; BD_TS_PACKET_SIZE];
|
||||
boundary[..self.remainder.len()].copy_from_slice(&self.remainder);
|
||||
boundary[self.remainder.len()..].copy_from_slice(&data[..need]);
|
||||
self.remainder.clear();
|
||||
self.process_packet(&boundary, &mut completed);
|
||||
offset = need;
|
||||
}
|
||||
|
||||
// Save leftover bytes for next call (cap at one packet to prevent unbounded growth)
|
||||
if offset < work.len() {
|
||||
let leftover = &work[offset..];
|
||||
// Aligned-packets fast path — reads directly out of `data`.
|
||||
while offset + BD_TS_PACKET_SIZE <= data.len() {
|
||||
let packet = &data[offset..offset + BD_TS_PACKET_SIZE];
|
||||
offset += BD_TS_PACKET_SIZE;
|
||||
self.process_packet(packet, &mut completed);
|
||||
}
|
||||
|
||||
// Save leftover bytes for next call (cap at one packet to
|
||||
// prevent unbounded growth on a desynchronised stream).
|
||||
if offset < data.len() {
|
||||
let leftover = &data[offset..];
|
||||
if leftover.len() < BD_TS_PACKET_SIZE {
|
||||
self.remainder.extend_from_slice(leftover);
|
||||
} else {
|
||||
// More than one full packet leftover — something is wrong, discard
|
||||
self.remainder.clear();
|
||||
}
|
||||
}
|
||||
@@ -222,6 +188,65 @@ impl TsDemuxer {
|
||||
completed
|
||||
}
|
||||
|
||||
/// Demux a single 192-byte BD-TS packet (4-byte TP_extra_header +
|
||||
/// 188-byte TS). Routes payload bytes into the per-PID
|
||||
/// `PesAssembler`; completed PES packets are pushed onto
|
||||
/// `completed` so the caller's allocation amortises across the
|
||||
/// batch.
|
||||
fn process_packet(&mut self, packet: &[u8], completed: &mut Vec<PesPacket>) {
|
||||
// Sync byte check skips malformed packets.
|
||||
if packet[4] != SYNC_BYTE {
|
||||
return;
|
||||
}
|
||||
let ts = &packet[4..]; // 188-byte standard TS packet
|
||||
|
||||
let pid = (((ts[1] & 0x1F) as u16) << 8) | ts[2] as u16;
|
||||
let pusi = ts[1] & 0x40 != 0; // Payload Unit Start Indicator
|
||||
let adaptation = (ts[3] >> 4) & 0x03;
|
||||
|
||||
let idx = if (pid as usize) < self.pid_index.len() {
|
||||
self.pid_index[pid as usize]
|
||||
} else {
|
||||
-1
|
||||
};
|
||||
if idx < 0 {
|
||||
return;
|
||||
}
|
||||
let asm = &mut self.assemblers[idx as usize];
|
||||
|
||||
let payload_start = if adaptation == 0x03 || adaptation == 0x02 {
|
||||
let af_len = ts[4] as usize;
|
||||
if af_len > 183 {
|
||||
return; // Malformed: AF length exceeds TS payload
|
||||
}
|
||||
5 + af_len
|
||||
} else {
|
||||
4
|
||||
};
|
||||
|
||||
if payload_start >= TS_PACKET_SIZE {
|
||||
return;
|
||||
}
|
||||
// adaptation == 0x02 → AF only, no payload.
|
||||
if adaptation == 0x02 {
|
||||
return;
|
||||
}
|
||||
|
||||
let payload = &ts[payload_start..];
|
||||
|
||||
if pusi {
|
||||
let (pts, dts, pes_data_start) = parse_pes_header(payload);
|
||||
if let Some(prev) = asm.start(pts, dts) {
|
||||
completed.push(prev);
|
||||
}
|
||||
if pes_data_start < payload.len() {
|
||||
asm.push(&payload[pes_data_start..]);
|
||||
}
|
||||
} else {
|
||||
asm.push(payload);
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush all assemblers, returning any remaining PES packets.
|
||||
pub fn flush(&mut self) -> Vec<PesPacket> {
|
||||
let mut completed = Vec::new();
|
||||
|
||||
+272
-35
@@ -42,28 +42,39 @@ impl<W: Write> TsMuxer<W> {
|
||||
/// Write a PES frame as BD-TS packets.
|
||||
/// Video frame data is expected as length-prefixed NALUs (MKV/PES format)
|
||||
/// and is converted to Annex B for transport stream.
|
||||
pub fn write_frame(&mut self, track: usize, pts_ns: i64, data: &[u8]) -> io::Result<()> {
|
||||
pub fn write_frame(
|
||||
&mut self,
|
||||
track: usize,
|
||||
pts_ns: i64,
|
||||
keyframe: bool,
|
||||
data: &[u8],
|
||||
) -> io::Result<()> {
|
||||
if track >= self.pids.len() {
|
||||
return Ok(()); // unknown track, skip
|
||||
}
|
||||
let base = *self.base_pts_ns.get_or_insert(pts_ns);
|
||||
let pts_ns = pts_ns - base;
|
||||
|
||||
let pid = self.pids[track];
|
||||
let is_video = (0x1011..=0x101F).contains(&pid);
|
||||
|
||||
// For video: convert length-prefixed NALUs to Annex B (start codes)
|
||||
// On first keyframe, prepend parameter sets from codec_private
|
||||
// Drop non-key video before any keyframe — decoder has no IDR or
|
||||
// parameter sets to anchor on.
|
||||
if is_video && !keyframe && !self.params_written[track] {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let base = *self.base_pts_ns.get_or_insert(pts_ns);
|
||||
let pts_ns = pts_ns - base;
|
||||
|
||||
// For video: convert length-prefixed NALUs to Annex B (start codes).
|
||||
// Prepend codec_private parameter sets on the FIRST keyframe only.
|
||||
let es_data = if is_video && !data.is_empty() {
|
||||
let mut annex_b = Vec::new();
|
||||
// Prepend codec_private parameter sets on first keyframe
|
||||
if !self.params_written[track] {
|
||||
if keyframe && !self.params_written[track] {
|
||||
if let Some(ref cp) = self.codec_privates[track] {
|
||||
if let Some(params) = hvcc_to_annex_b(cp) {
|
||||
annex_b.extend_from_slice(¶ms);
|
||||
self.params_written[track] = true;
|
||||
}
|
||||
}
|
||||
self.params_written[track] = true;
|
||||
}
|
||||
annex_b.extend_from_slice(&length_prefixed_to_annex_b(data));
|
||||
annex_b
|
||||
@@ -85,8 +96,24 @@ impl<W: Write> TsMuxer<W> {
|
||||
let mut first = true;
|
||||
while offset < pes_packet.len() {
|
||||
let remaining = pes_packet.len() - offset;
|
||||
let payload_len = remaining.min(TS_PAYLOAD);
|
||||
let need_stuffing = payload_len < TS_PAYLOAD;
|
||||
|
||||
// Invariant: TP_extra(4) + TS_header(4) + AF(af_bytes) + payload(payload_len) = 192,
|
||||
// i.e. af_bytes + payload_len = TS_PAYLOAD (184).
|
||||
// RAI on first packet of a keyframe video PES requires AF with flags=0x40.
|
||||
let want_rai = first && keyframe && is_video;
|
||||
|
||||
// Pick payload_len and af_bytes per case.
|
||||
let (af_bytes, payload_len): (usize, usize) = if want_rai {
|
||||
// Minimum AF = 2 bytes (length=1, flags=0x40). Payload caps at 182.
|
||||
let max_payload = TS_PAYLOAD - 2;
|
||||
let p = remaining.min(max_payload);
|
||||
(TS_PAYLOAD - p, p)
|
||||
} else if remaining >= TS_PAYLOAD {
|
||||
(0, TS_PAYLOAD) // no AF, full payload
|
||||
} else {
|
||||
// Stuffing-only AF, payload = remaining.
|
||||
(TS_PAYLOAD - remaining, remaining)
|
||||
};
|
||||
|
||||
// TP_extra_header (4 bytes — arrival time, set to 0)
|
||||
let tp_extra = [0u8; 4];
|
||||
@@ -102,38 +129,45 @@ impl<W: Write> TsMuxer<W> {
|
||||
ts_header[1] |= 0x40; // PUSI
|
||||
}
|
||||
ts_header[2] = pid as u8;
|
||||
ts_header[3] = 0x10 | cc; // no adaptation, has payload
|
||||
ts_header[3] = if af_bytes > 0 {
|
||||
0x30 | cc // AF + payload
|
||||
} else {
|
||||
0x10 | cc // payload only
|
||||
};
|
||||
|
||||
if need_stuffing {
|
||||
// Adaptation field for stuffing
|
||||
let stuff_len = TS_PAYLOAD - payload_len;
|
||||
ts_header[3] = 0x30 | cc; // adaptation + payload
|
||||
self.writer.write_all(&tp_extra)?;
|
||||
self.writer.write_all(&ts_header)?;
|
||||
|
||||
self.writer.write_all(&tp_extra)?;
|
||||
self.writer.write_all(&ts_header)?;
|
||||
|
||||
// Write adaptation field: length byte + flags byte + 0xFF padding
|
||||
// stuff_len == 1: AF length = 0 (just the length byte, no flags)
|
||||
// stuff_len >= 2: AF length = stuff_len-1, flags = 0, rest 0xFF
|
||||
if af_bytes > 0 {
|
||||
static STUFF_FF: [u8; 184] = [0xFF; 184];
|
||||
if stuff_len == 1 {
|
||||
self.writer.write_all(&[0u8])?; // adaptation_field_length = 0
|
||||
if want_rai {
|
||||
// RAI AF: length byte + flags(0x40) + (af_bytes - 2) stuffing.
|
||||
let af_len_field = (af_bytes - 1) as u8;
|
||||
self.writer.write_all(&[af_len_field])?;
|
||||
self.writer.write_all(&[0x40u8])?;
|
||||
let stuff = af_bytes - 2;
|
||||
if stuff > 0 {
|
||||
self.writer.write_all(&STUFF_FF[..stuff])?;
|
||||
}
|
||||
} else {
|
||||
self.writer.write_all(&[(stuff_len - 1) as u8])?; // AF length
|
||||
self.writer.write_all(&[0u8])?; // flags
|
||||
if stuff_len > 2 {
|
||||
self.writer.write_all(&STUFF_FF[..stuff_len - 2])?;
|
||||
// Stuffing-only AF.
|
||||
// af_bytes == 1: length=0, no flags.
|
||||
// af_bytes >= 2: length = af_bytes-1, flags=0, rest 0xFF.
|
||||
if af_bytes == 1 {
|
||||
self.writer.write_all(&[0u8])?;
|
||||
} else {
|
||||
self.writer.write_all(&[(af_bytes - 1) as u8])?;
|
||||
self.writer.write_all(&[0u8])?;
|
||||
if af_bytes > 2 {
|
||||
self.writer.write_all(&STUFF_FF[..af_bytes - 2])?;
|
||||
}
|
||||
}
|
||||
}
|
||||
self.writer
|
||||
.write_all(&pes_packet[offset..offset + payload_len])?;
|
||||
} else {
|
||||
self.writer.write_all(&tp_extra)?;
|
||||
self.writer.write_all(&ts_header)?;
|
||||
self.writer
|
||||
.write_all(&pes_packet[offset..offset + payload_len])?;
|
||||
}
|
||||
|
||||
self.writer
|
||||
.write_all(&pes_packet[offset..offset + payload_len])?;
|
||||
|
||||
offset += payload_len;
|
||||
first = false;
|
||||
}
|
||||
@@ -257,3 +291,206 @@ fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const BD_PACKET_SIZE: usize = 192;
|
||||
const VIDEO_PID: u16 = 0x1011;
|
||||
|
||||
/// Parsed BD-TS packet (192 bytes total: 4 TP_extra + 4 TS header + 184 body).
|
||||
struct TsPacket {
|
||||
pid: u16,
|
||||
pusi: bool,
|
||||
#[allow(dead_code)]
|
||||
cc: u8,
|
||||
/// Adaptation field body (length byte stripped) when present.
|
||||
af: Option<Vec<u8>>,
|
||||
/// Payload bytes (after AF, if any).
|
||||
payload: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Walk 192-byte BD-TS packets.
|
||||
fn parse_bd_ts(buf: &[u8]) -> Vec<TsPacket> {
|
||||
let mut out = Vec::new();
|
||||
for chunk in buf.chunks(BD_PACKET_SIZE) {
|
||||
if chunk.len() != BD_PACKET_SIZE {
|
||||
break;
|
||||
}
|
||||
// Skip TP_extra_header (4 bytes), parse TS header.
|
||||
let h = &chunk[4..];
|
||||
assert_eq!(h[0], 0x47, "bad sync byte");
|
||||
let pusi = (h[1] & 0x40) != 0;
|
||||
let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16;
|
||||
let afc = (h[3] >> 4) & 0x03;
|
||||
let cc = h[3] & 0x0F;
|
||||
let body = &h[4..]; // 184 bytes
|
||||
|
||||
let (af, payload) = match afc {
|
||||
0b01 => (None, body.to_vec()),
|
||||
0b11 => {
|
||||
let af_len = body[0] as usize;
|
||||
let af_body = body[1..1 + af_len].to_vec();
|
||||
let payload = body[1 + af_len..].to_vec();
|
||||
(Some(af_body), payload)
|
||||
}
|
||||
0b10 => {
|
||||
let af_len = body[0] as usize;
|
||||
(Some(body[1..1 + af_len].to_vec()), Vec::new())
|
||||
}
|
||||
_ => (None, Vec::new()),
|
||||
};
|
||||
|
||||
out.push(TsPacket {
|
||||
pid,
|
||||
pusi,
|
||||
cc,
|
||||
af,
|
||||
payload,
|
||||
});
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Build a fake HEVC NAL with a 4-byte length prefix.
|
||||
/// nal_type=19/20 are IDR; 1 is non-key (TRAIL_N/R).
|
||||
fn fake_hevc_nal(nal_type: u8, body_len: usize) -> Vec<u8> {
|
||||
let mut nal = Vec::with_capacity(2 + body_len);
|
||||
// 2-byte NAL header: forbidden_zero(1)=0 | nal_unit_type(6) | layer_id(6)=0 | tid_plus1(3)=1
|
||||
nal.push((nal_type & 0x3F) << 1);
|
||||
nal.push(0x01);
|
||||
for i in 0..body_len {
|
||||
nal.push((i & 0xFF) as u8);
|
||||
}
|
||||
let mut framed = Vec::with_capacity(4 + nal.len());
|
||||
framed.extend_from_slice(&(nal.len() as u32).to_be_bytes());
|
||||
framed.extend_from_slice(&nal);
|
||||
framed
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keyframe_param_threads_through() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let idr = fake_hevc_nal(19, 100);
|
||||
mux.write_frame(0, 0, true, &idr).unwrap();
|
||||
let p = fake_hevc_nal(1, 80);
|
||||
mux.write_frame(0, 41_000_000, false, &p).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
assert!(!sink.is_empty());
|
||||
let packets = parse_bd_ts(&sink);
|
||||
assert!(packets.iter().any(|p| p.pid == VIDEO_PID && p.pusi));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rai_set_on_first_packet_of_keyframe_pes() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let idr = fake_hevc_nal(19, 200);
|
||||
mux.write_frame(0, 0, true, &idr).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
let packets = parse_bd_ts(&sink);
|
||||
let first_pusi = packets
|
||||
.iter()
|
||||
.find(|p| p.pid == VIDEO_PID && p.pusi)
|
||||
.expect("video PUSI packet exists");
|
||||
let af = first_pusi.af.as_ref().expect("AF present on keyframe PES");
|
||||
assert!(!af.is_empty(), "AF body has flags byte");
|
||||
assert_eq!(af[0] & 0x40, 0x40, "RAI bit set");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rai_clear_on_non_keyframe_pes() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let idr = fake_hevc_nal(19, 100);
|
||||
mux.write_frame(0, 0, true, &idr).unwrap();
|
||||
let p = fake_hevc_nal(1, 100);
|
||||
mux.write_frame(0, 41_000_000, false, &p).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
let packets = parse_bd_ts(&sink);
|
||||
// Second PUSI packet on the video PID belongs to the non-key frame.
|
||||
let pusi_video: Vec<&TsPacket> = packets
|
||||
.iter()
|
||||
.filter(|p| p.pid == VIDEO_PID && p.pusi)
|
||||
.collect();
|
||||
assert!(pusi_video.len() >= 2, "two PUSI packets expected");
|
||||
let second = pusi_video[1];
|
||||
match &second.af {
|
||||
None => {}
|
||||
Some(af) if af.is_empty() => {} // length=0 case
|
||||
Some(af) => assert_eq!(af[0] & 0x40, 0, "RAI must be clear on non-key PES"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codec_private_prepended_only_on_first_keyframe() {
|
||||
// Build a minimal hvcC with one recognizable NAL.
|
||||
let marker: &[u8] = &[0xDE, 0xAD, 0xBE, 0xEF, 0xCA, 0xFE];
|
||||
let mut hvcc = vec![0u8; 22];
|
||||
hvcc.push(1); // numArrays
|
||||
hvcc.push(32); // VPS NAL type byte (high bits arbitrary)
|
||||
hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
|
||||
hvcc.extend_from_slice(&(marker.len() as u16).to_be_bytes());
|
||||
hvcc.extend_from_slice(marker);
|
||||
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
mux.set_codec_private(0, hvcc);
|
||||
// Non-IDR before any IDR: should be dropped.
|
||||
let p = fake_hevc_nal(1, 50);
|
||||
mux.write_frame(0, 0, false, &p).unwrap();
|
||||
// IDR: should carry codec_private NALs prepended.
|
||||
let idr = fake_hevc_nal(19, 50);
|
||||
mux.write_frame(0, 41_000_000, true, &idr).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
let packets = parse_bd_ts(&sink);
|
||||
// Concatenate all video PID payloads in emission order.
|
||||
let video_bytes: Vec<u8> = packets
|
||||
.iter()
|
||||
.filter(|p| p.pid == VIDEO_PID)
|
||||
.flat_map(|p| p.payload.clone())
|
||||
.collect();
|
||||
// marker bytes must appear in the stream (codec_private was prepended).
|
||||
let pos_marker = video_bytes
|
||||
.windows(marker.len())
|
||||
.position(|w| w == marker)
|
||||
.expect("codec_private marker bytes present in TS payload");
|
||||
// Find IDR body byte (0x26 = (19<<1)). pos_idr must be AFTER marker.
|
||||
let idr_header = (19u8 << 1) & 0x7E;
|
||||
let pos_idr = video_bytes
|
||||
.iter()
|
||||
.position(|&b| b == idr_header)
|
||||
.expect("IDR NAL header present in TS payload");
|
||||
assert!(
|
||||
pos_marker < pos_idr,
|
||||
"codec_private must precede IDR in TS payload"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn non_key_before_first_keyframe_dropped() {
|
||||
let mut sink: Vec<u8> = Vec::new();
|
||||
{
|
||||
let mut mux = TsMuxer::new(&mut sink, &[VIDEO_PID]);
|
||||
let p = fake_hevc_nal(1, 80);
|
||||
mux.write_frame(0, 0, false, &p).unwrap();
|
||||
mux.finish().unwrap();
|
||||
}
|
||||
// Nothing should be emitted for that PID.
|
||||
let packets = parse_bd_ts(&sink);
|
||||
assert!(
|
||||
!packets.iter().any(|p| p.pid == VIDEO_PID),
|
||||
"non-key before first keyframe must be dropped"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
+17
@@ -17,6 +17,10 @@ pub struct PesFrame {
|
||||
pub keyframe: bool,
|
||||
/// Raw elementary stream data (NAL units, audio samples, etc).
|
||||
pub data: Vec<u8>,
|
||||
/// Optional duration in nanoseconds. In-memory only; not part of
|
||||
/// the on-wire serialization. Currently set by the PGS parser so
|
||||
/// the MKV muxer can emit `BlockDuration`.
|
||||
pub duration_ns: Option<u64>,
|
||||
}
|
||||
|
||||
impl PesFrame {
|
||||
@@ -64,6 +68,7 @@ impl PesFrame {
|
||||
pts,
|
||||
keyframe,
|
||||
data,
|
||||
duration_ns: None,
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -74,6 +79,7 @@ impl PesFrame {
|
||||
pts: frame.pts_ns,
|
||||
keyframe: frame.keyframe,
|
||||
data: frame.data,
|
||||
duration_ns: frame.duration_ns,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -115,6 +121,16 @@ pub trait Stream: Send {
|
||||
fn headers_ready(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
/// Cumulative count of read errors the stream skipped past (e.g.
|
||||
/// zero-filled bad sectors on a live drive). Default `0` for
|
||||
/// streams that don't have a notion of skip-on-error (file ISO,
|
||||
/// network, stdio, the pipeline highway, etc.); concrete impls
|
||||
/// with adaptive retry (`DiscStream` on the drive single-pass
|
||||
/// path) override.
|
||||
fn errors(&self) -> u64 {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps any output stream and counts bytes written.
|
||||
@@ -183,6 +199,7 @@ mod tests {
|
||||
pts,
|
||||
keyframe: track == 0 && pts == 0,
|
||||
data: vec![track as u8, (pts & 0xff) as u8, 0xAA],
|
||||
duration_ns: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
//! Linux `statfs64`-based filesystem type detection.
|
||||
//!
|
||||
//! Recognised local-FS magics: ext2/3/4, xfs, btrfs, tmpfs. NFS is the
|
||||
//! one network FS this layer cares about (the buffering decision keys
|
||||
//! off it). Anything else maps to [`FsType::Unknown`].
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::unix::ffi::OsStrExt;
|
||||
use std::os::unix::io::RawFd;
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
// Magic numbers from `<linux/magic.h>`. Kept literal here so we don't
|
||||
// depend on libc exposing each one — only `NFS_SUPER_MAGIC` is
|
||||
// guaranteed to be present across libc / musl revisions.
|
||||
const EXT2_SUPER_MAGIC: i64 = 0xEF53;
|
||||
const XFS_SUPER_MAGIC: i64 = 0x5846_5342;
|
||||
const BTRFS_SUPER_MAGIC: i64 = 0x9123_683E;
|
||||
const TMPFS_MAGIC: i64 = 0x0102_1994;
|
||||
|
||||
/// Classify an `f_type` magic from `statfs`/`fstatfs`. Single source
|
||||
/// of truth for the magic comparisons used by both the path-based
|
||||
/// (`detect_impl`) and fd-based (`detect_fd_impl`) entry points.
|
||||
///
|
||||
/// Cast to `i64` because `statfs::f_type` is signed `__fsword_t` on
|
||||
/// glibc and unsigned `c_ulong` on musl; a portable comparison needs a
|
||||
/// common type. On glibc x86_64 both already are i64 — clippy flags
|
||||
/// the cast as unnecessary on that target only, but we need it for
|
||||
/// musl.
|
||||
#[allow(clippy::unnecessary_cast)]
|
||||
fn classify_f_type(f_type: i64) -> FsType {
|
||||
let nfs_magic = libc::NFS_SUPER_MAGIC as i64;
|
||||
if f_type == nfs_magic {
|
||||
return FsType::Nfs;
|
||||
}
|
||||
match f_type {
|
||||
EXT2_SUPER_MAGIC | XFS_SUPER_MAGIC | BTRFS_SUPER_MAGIC | TMPFS_MAGIC => FsType::Local,
|
||||
_ => FsType::Unknown,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn detect_impl(path: &Path) -> FsType {
|
||||
let cpath = match CString::new(path.as_os_str().as_bytes()) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return FsType::Unknown,
|
||||
};
|
||||
// `statfs64` is repr(C); zeroing is the documented init pattern for
|
||||
// the kernel uapi struct.
|
||||
let mut buf: libc::statfs64 = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::statfs64(cpath.as_ptr(), &mut buf) };
|
||||
if rc != 0 {
|
||||
return FsType::Unknown;
|
||||
}
|
||||
#[allow(clippy::unnecessary_cast)]
|
||||
classify_f_type(buf.f_type as i64)
|
||||
}
|
||||
|
||||
/// fd-based variant of [`detect_impl`]. Same return semantics. Used
|
||||
/// by the writeback pipeline, which knows the open `File` but not its
|
||||
/// original path.
|
||||
pub(super) fn detect_fd_impl(fd: RawFd) -> FsType {
|
||||
let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::fstatfs(fd, &mut buf) };
|
||||
if rc != 0 {
|
||||
return FsType::Unknown;
|
||||
}
|
||||
#[allow(clippy::unnecessary_cast)]
|
||||
classify_f_type(buf.f_type as i64)
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
//! macOS `statfs`-based filesystem type detection.
|
||||
//!
|
||||
//! macOS exposes a textual `f_fstypename` field (e.g. `"apfs"`, `"hfs"`,
|
||||
//! `"nfs"`, `"smbfs"`) on its `statfs` struct, which is far more
|
||||
//! reliable than chasing magic numbers. Anything starting with `"nfs"`
|
||||
//! counts as NFS; anything else recognised maps to `Local`; otherwise
|
||||
//! `Unknown`.
|
||||
|
||||
use std::ffi::CString;
|
||||
use std::os::unix::ffi::OsStrExt;
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
pub(super) fn detect_impl(path: &Path) -> FsType {
|
||||
let cpath = match CString::new(path.as_os_str().as_bytes()) {
|
||||
Ok(c) => c,
|
||||
Err(_) => return FsType::Unknown,
|
||||
};
|
||||
let mut buf: libc::statfs = unsafe { std::mem::zeroed() };
|
||||
let rc = unsafe { libc::statfs(cpath.as_ptr(), &mut buf) };
|
||||
if rc != 0 {
|
||||
return FsType::Unknown;
|
||||
}
|
||||
// `f_fstypename` is a NUL-terminated C string of length MFSTYPENAMELEN.
|
||||
// SAFETY: libc guarantees the field is initialised to a NUL-terminated
|
||||
// string by a successful statfs.
|
||||
let name_ptr = buf.f_fstypename.as_ptr();
|
||||
let cstr = unsafe { std::ffi::CStr::from_ptr(name_ptr) };
|
||||
let name = cstr.to_bytes();
|
||||
if name.starts_with(b"nfs") {
|
||||
return FsType::Nfs;
|
||||
}
|
||||
// Recognised local types. SMB is not NFS but the buffering-policy
|
||||
// outcome on macOS is the same as for any other local FS — there's
|
||||
// no `WritebackFile` machinery to opt out of on this OS.
|
||||
match name {
|
||||
b"apfs" | b"hfs" | b"exfat" | b"msdos" | b"tmpfs" | b"smbfs" | b"webdav" => FsType::Local,
|
||||
_ => FsType::Unknown,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
//! Filesystem-type detection.
|
||||
//!
|
||||
//! The buffering architecture (Phase 2) selects different output sinks
|
||||
//! for local vs network filesystems: NFS gets the adaptive
|
||||
//! `WritebackFile` machinery on Linux; local disks get `LocalFileSink`
|
||||
//! and rely on the kernel's default writeback policy. This module
|
||||
//! provides the construction-site primitive that picks which one.
|
||||
//!
|
||||
//! Per the per-OS file-split convention, the actual `statfs` call lives
|
||||
//! in the matching platform file (`linux.rs`, `macos.rs`, `windows.rs`,
|
||||
//! `other.rs`); this `mod.rs` exposes only the cross-platform enum and
|
||||
//! the `detect` entry point.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
/// What kind of filesystem a path lives on, to the extent we can tell
|
||||
/// cheaply at construction time.
|
||||
///
|
||||
/// `Unknown` is the fail-open default: a misdetection here should not
|
||||
/// be load-bearing for correctness, only for the choice of sink (and
|
||||
/// hence buffering policy). Callers that need a binary local/non-local
|
||||
/// answer should treat `Unknown` as "probably local".
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum FsType {
|
||||
/// A local on-disk filesystem (ext4, xfs, btrfs, apfs, ntfs, …).
|
||||
Local,
|
||||
/// A network filesystem with NFS-like semantics. The current
|
||||
/// detector lumps SMB / UNC into this on Windows because the
|
||||
/// buffering policy outcome is the same.
|
||||
Nfs,
|
||||
/// `statfs` failed, the filesystem type is not on our recognised
|
||||
/// list, or we are on an OS without a real implementation.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
mod linux;
|
||||
#[cfg(target_os = "macos")]
|
||||
mod macos;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
mod other;
|
||||
#[cfg(target_os = "windows")]
|
||||
mod windows;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux::detect_impl;
|
||||
#[cfg(target_os = "macos")]
|
||||
use macos::detect_impl;
|
||||
#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
|
||||
use other::detect_impl;
|
||||
#[cfg(target_os = "windows")]
|
||||
use windows::detect_impl;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use linux::detect_fd_impl;
|
||||
|
||||
/// Best-effort classification of the filesystem under `path`.
|
||||
///
|
||||
/// Falls back to [`FsType::Unknown`] on any syscall error or unrecognised
|
||||
/// filesystem signature. Never panics. Never blocks beyond the cost of
|
||||
/// a single `statfs(2)` (Unix) or a string check (Windows).
|
||||
pub fn detect(path: &Path) -> FsType {
|
||||
detect_impl(path)
|
||||
}
|
||||
|
||||
/// fd-based classification. Same return semantics as [`detect`], but
|
||||
/// takes a `RawFd` so callers that only have an open file (notably
|
||||
/// [`crate::io::writeback::WritebackPipeline`]) don't have to
|
||||
/// round-trip through the path.
|
||||
///
|
||||
/// Only implemented on Linux; other platforms return
|
||||
/// [`FsType::Unknown`] (none of them have a writeback policy that
|
||||
/// keys off this classification today).
|
||||
#[cfg(target_os = "linux")]
|
||||
pub fn detect_fd(fd: std::os::unix::io::RawFd) -> FsType {
|
||||
detect_fd_impl(fd)
|
||||
}
|
||||
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
#[allow(dead_code)] // API parity with the linux impl; callers cfg-gate.
|
||||
pub fn detect_fd(_fd: i32) -> FsType {
|
||||
FsType::Unknown
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn detect_does_not_panic_on_missing_path() {
|
||||
// Non-existent path should fall through to Unknown, not panic.
|
||||
let p = std::path::Path::new("/this/path/should/not/exist/freemkv-test");
|
||||
let r = detect(p);
|
||||
// We don't assert == Unknown because Windows' heuristic looks at
|
||||
// the leading bytes and might still classify; just confirm the
|
||||
// call returns rather than panicking.
|
||||
let _ = r;
|
||||
}
|
||||
|
||||
#[cfg(target_os = "macos")]
|
||||
#[test]
|
||||
fn macos_tmp_is_local() {
|
||||
// `/tmp` on macOS dev rigs is APFS via the symlink to
|
||||
// `/private/tmp`. Either way, never NFS.
|
||||
let r = detect(std::path::Path::new("/tmp"));
|
||||
assert!(
|
||||
matches!(r, FsType::Local | FsType::Unknown),
|
||||
"expected Local or Unknown for /tmp on macOS, got {r:?}",
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[test]
|
||||
fn linux_tmp_is_local_or_unknown() {
|
||||
// `/tmp` is tmpfs on most distros (which we recognise) but
|
||||
// could be ext4 on others. NFS would be unusual.
|
||||
let r = detect(std::path::Path::new("/tmp"));
|
||||
assert!(
|
||||
matches!(r, FsType::Local | FsType::Unknown),
|
||||
"expected Local or Unknown for /tmp on Linux, got {r:?}",
|
||||
);
|
||||
}
|
||||
|
||||
/// Real NFS exercise needs an actual NFS mount and isn't available
|
||||
/// in CI. Kept here as a manual probe.
|
||||
#[test]
|
||||
#[ignore = "needs an NFS mount path (e.g. /mnt/nfs/...) to validate live"]
|
||||
fn nfs_path_classified_as_nfs() {
|
||||
// Operator passes the mount as FREEMKV_NFS_PROBE; gated behind
|
||||
// `--ignored` because there's no portable NFS path.
|
||||
let p = std::env::var("FREEMKV_NFS_PROBE").expect("set FREEMKV_NFS_PROBE");
|
||||
assert_eq!(detect(std::path::Path::new(&p)), FsType::Nfs);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
//! Fallback fs-type detection for platforms without a specific impl.
|
||||
//!
|
||||
//! Always returns [`FsType::Unknown`]. Callers treat that as
|
||||
//! "probably local" and pick `LocalFileSink`.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
pub(super) fn detect_impl(_path: &Path) -> FsType {
|
||||
FsType::Unknown
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
//! Windows filesystem-type detection.
|
||||
//!
|
||||
//! Heuristic-only: any UNC path (`\\server\share\...`) is treated as a
|
||||
//! network mount and bucketed into `Nfs`. Strictly, SMB is not NFS, but
|
||||
//! the buffering-policy outcome for our purposes is the same — there is
|
||||
//! no platform `WritebackFile` machinery on Windows yet, so the worst
|
||||
//! case of a false positive is using `LocalFileSink` regardless. A
|
||||
//! proper `GetVolumeInformation` query is a Phase 4 concern.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
use super::FsType;
|
||||
|
||||
pub(super) fn detect_impl(path: &Path) -> FsType {
|
||||
// `Path::starts_with("\\\\")` won't match because component-wise
|
||||
// matching strips the prefix. Look at the raw OsStr instead.
|
||||
let s = path.as_os_str();
|
||||
// OsStr -> [u16] is the platform-correct way on Windows, but
|
||||
// checking the leading bytes via `to_string_lossy` is good enough
|
||||
// for a UNC prefix probe and works on any encoding.
|
||||
let lossy = s.to_string_lossy();
|
||||
if lossy.starts_with("\\\\") || lossy.starts_with("//") {
|
||||
return FsType::Nfs;
|
||||
}
|
||||
FsType::Local
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
//! Platform-specific drive initialization and disc probing.
|
||||
|
||||
pub mod fs_type;
|
||||
pub mod mt1959;
|
||||
|
||||
use crate::error::Result;
|
||||
|
||||
+12
-146
@@ -1,148 +1,19 @@
|
||||
//! File-backed sector I/O — read and write 2048-byte sectors against
|
||||
//! an ISO image on disk.
|
||||
//! File-backed sector sink — write 2048-byte sectors to an ISO image
|
||||
//! on disk.
|
||||
//!
|
||||
//! [`FileSectorSource`] is the read side (open-only). [`FileSectorSink`]
|
||||
//! is the write side (create or open-rw); writes go through
|
||||
//! [`crate::io::WritebackFile`] so big sequential ISO writes share
|
||||
//! the same bounded-cache writeback pipeline used by sweep / patch /
|
||||
//! mux.
|
||||
//! The read-side counterpart ([`crate::io::file_sector_source::FileSectorSource`])
|
||||
//! lives under `io/` because its internals (read-ahead buffer, per-OS
|
||||
//! `fadvise`/`F_RDADVISE` hints) are I/O infrastructure rather than
|
||||
//! sector-trait business logic. Both types remain re-exported at
|
||||
//! [`crate::sector`] for ergonomic imports.
|
||||
|
||||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{Read, Seek, SeekFrom, Write};
|
||||
use std::fs::OpenOptions;
|
||||
use std::io::{Seek, SeekFrom, Write};
|
||||
use std::path::Path;
|
||||
|
||||
use crate::error::{Error, Result};
|
||||
|
||||
use super::{SectorSink, SectorSource};
|
||||
|
||||
/// Bytes-read threshold per `posix_fadvise(DONTNEED)` drop on the
|
||||
/// read side. Mirrors the writeback chunk size so the read-side
|
||||
/// page cache stays bounded the same way the write side does.
|
||||
#[cfg(target_os = "linux")]
|
||||
const READ_DROP_CHUNK_BYTES: u64 = 32 * 1024 * 1024;
|
||||
|
||||
/// SectorSource backed by a file (ISO image).
|
||||
///
|
||||
/// Seeks to `lba * 2048`, reads `count * 2048` bytes per call. The
|
||||
/// file is held directly: every `read_sectors` call performs an
|
||||
/// absolute seek, so a wrapping `BufReader` would have its buffer
|
||||
/// invalidated on every call (its internal cursor moves with the
|
||||
/// `Seek` impl) — pure overhead. Callers that benefit from buffered
|
||||
/// reads should compose their own `BufReader` at the `read_sectors`
|
||||
/// granularity they care about.
|
||||
pub struct FileSectorSource {
|
||||
file: File,
|
||||
capacity: u32,
|
||||
/// Bytes read since the last `posix_fadvise(DONTNEED)` drop.
|
||||
/// Only updated on Linux; on other targets it stays at 0.
|
||||
#[cfg(target_os = "linux")]
|
||||
bytes_read_since_drop: u64,
|
||||
/// Byte offset at which the current drop window starts (the
|
||||
/// next `posix_fadvise(DONTNEED)` call drops from here).
|
||||
#[cfg(target_os = "linux")]
|
||||
drop_window_start: u64,
|
||||
}
|
||||
|
||||
impl FileSectorSource {
|
||||
/// Open an existing ISO file for reading. Capacity is derived
|
||||
/// from `metadata().len() / 2048`. Returns
|
||||
/// [`Error::IsoTooLarge`] if the file would exceed the 32-bit
|
||||
/// LBA address space (~8 TB).
|
||||
pub fn open(path: &Path) -> std::io::Result<Self> {
|
||||
let file = File::open(path)?;
|
||||
let len = file.metadata()?.len();
|
||||
let sectors = len / 2048;
|
||||
if sectors > u32::MAX as u64 {
|
||||
return Err(Error::IsoTooLarge {
|
||||
path: path.to_string_lossy().into_owned(),
|
||||
}
|
||||
.into());
|
||||
}
|
||||
let capacity = sectors as u32;
|
||||
|
||||
// Hint sequential access on Linux so the kernel's readahead
|
||||
// window widens for the ISO sweep. Best-effort: return value
|
||||
// is ignored. On macOS / Windows this is a no-op.
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
use std::os::unix::io::AsRawFd;
|
||||
unsafe {
|
||||
libc::posix_fadvise(file.as_raw_fd(), 0, 0, libc::POSIX_FADV_SEQUENTIAL);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
file,
|
||||
capacity,
|
||||
#[cfg(target_os = "linux")]
|
||||
bytes_read_since_drop: 0,
|
||||
#[cfg(target_os = "linux")]
|
||||
drop_window_start: 0,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Implement the legacy `SectorSource` trait. The blanket impl in
|
||||
// `super` produces the `SectorSource` impl automatically — no need
|
||||
// to write both, and writing both would conflict. This keeps the
|
||||
// 0.17 method-resolution path intact (callers with `SectorSource`
|
||||
// in scope can still write `fsr.read_sectors(..)` against a
|
||||
// `FileSectorSource`).
|
||||
impl SectorSource for FileSectorSource {
|
||||
fn capacity_sectors(&self) -> u32 {
|
||||
self.capacity
|
||||
}
|
||||
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
lba: u32,
|
||||
count: u16,
|
||||
buf: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
let offset = lba as u64 * 2048;
|
||||
let bytes = count as usize * 2048;
|
||||
self.file
|
||||
.seek(SeekFrom::Start(offset))
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
self.file
|
||||
.read_exact(&mut buf[..bytes])
|
||||
.map_err(|e| Error::IoError { source: e })?;
|
||||
|
||||
// On Linux, periodically drop the just-read region from the
|
||||
// page cache to keep cache pressure bounded during multi-GB
|
||||
// sequential ISO reads. Mirrors the write-side pipeline.
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
use std::os::unix::io::AsRawFd;
|
||||
self.bytes_read_since_drop += bytes as u64;
|
||||
if self.bytes_read_since_drop >= READ_DROP_CHUNK_BYTES {
|
||||
let drop_start = self.drop_window_start;
|
||||
let drop_len = self.bytes_read_since_drop;
|
||||
let t0 = std::time::Instant::now();
|
||||
unsafe {
|
||||
libc::posix_fadvise(
|
||||
self.file.as_raw_fd(),
|
||||
drop_start as i64,
|
||||
drop_len as i64,
|
||||
libc::POSIX_FADV_DONTNEED,
|
||||
);
|
||||
}
|
||||
let elapsed_ms = t0.elapsed().as_millis();
|
||||
let start_lba = drop_start / 2048;
|
||||
let end_lba = (drop_start + drop_len) / 2048;
|
||||
tracing::trace!(
|
||||
target: "mux",
|
||||
"FileSectorSource fadvise DONTNEED lba=[{start_lba}..{end_lba}) bytes={drop_len} elapsed_ms={elapsed_ms}"
|
||||
);
|
||||
self.drop_window_start = drop_start + drop_len;
|
||||
self.bytes_read_since_drop = 0;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(bytes)
|
||||
}
|
||||
}
|
||||
use super::SectorSink;
|
||||
|
||||
/// SectorSink backed by a file (ISO image).
|
||||
///
|
||||
@@ -211,13 +82,8 @@ impl SectorSink for FileSectorSink {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
// Bring the 0.18 trait into scope (not super::*: the super
|
||||
// module also re-exports the legacy `SectorSource`, and
|
||||
// having both `SectorSource::read_sectors` and
|
||||
// `SectorSource::read_sectors` visible would force every
|
||||
// call site to disambiguate). External consumers see the
|
||||
// same surface this test exercises.
|
||||
use super::{FileSectorSink, FileSectorSource};
|
||||
use super::FileSectorSink;
|
||||
use crate::io::file_sector_source::FileSectorSource;
|
||||
use crate::sector::{SectorSink, SectorSource};
|
||||
use tempfile::tempdir;
|
||||
|
||||
|
||||
+5
-2
@@ -6,7 +6,7 @@
|
||||
//! ISO images are opened for read OR write at construction time.
|
||||
//!
|
||||
//! - [`SectorSource`] is implemented by `Drive` (hardware) and
|
||||
//! [`FileSectorSource`] / `IsoSectorReader` (file-backed).
|
||||
//! [`FileSectorSource`] (file-backed).
|
||||
//! - [`SectorSink`] is implemented by [`FileSectorSink`]
|
||||
//! (ISO-backed) and sweep/patch consumer adapters.
|
||||
//! - [`DecryptingSectorSource`] is a decorator that wraps any
|
||||
@@ -15,6 +15,7 @@
|
||||
|
||||
pub mod decrypting;
|
||||
pub mod file;
|
||||
pub mod prefetched;
|
||||
|
||||
use crate::error::Result;
|
||||
|
||||
@@ -109,5 +110,7 @@ pub trait SectorSink: Send {
|
||||
fn finish(self: Box<Self>) -> Result<()>;
|
||||
}
|
||||
|
||||
pub use crate::io::file_sector_source::FileSectorSource;
|
||||
pub use decrypting::DecryptingSectorSource;
|
||||
pub use file::{FileSectorSink, FileSectorSource};
|
||||
pub use file::FileSectorSink;
|
||||
pub use prefetched::PrefetchedSectorSource;
|
||||
|
||||
@@ -0,0 +1,285 @@
|
||||
//! `PrefetchedSectorSource` — runs the wrapped read+decrypt in a
|
||||
//! dedicated producer thread and surfaces the prepared plaintext
|
||||
//! buffers on demand via a bounded channel.
|
||||
//!
|
||||
//! ## Why
|
||||
//!
|
||||
//! The mux consumer (demux + codec parsing + frame output) is
|
||||
//! single-threaded by nature (streams are sequential). The mux
|
||||
//! producer (read sectors + AACS decrypt) is also single-threaded
|
||||
//! per-call but does CPU-heavy work (AES per 6144-byte unit). Running
|
||||
//! both on the same thread means the disk and decrypt cores sit idle
|
||||
//! while the demux runs, and vice versa.
|
||||
//!
|
||||
//! Splitting them across two threads with a bounded channel between
|
||||
//! lets both run in parallel — peak throughput becomes
|
||||
//! `min(producer_rate, consumer_rate)` instead of
|
||||
//! `1 / (1/producer + 1/consumer)`.
|
||||
//!
|
||||
//! ## Lifecycle
|
||||
//!
|
||||
//! The producer thread is spawned by [`PrefetchedSectorSource::new`].
|
||||
//! It walks the supplied extent list in order, reads the configured
|
||||
//! batch size at each LBA, and sends the resulting plaintext buffer
|
||||
//! into a [`crossbeam_channel::bounded`] channel of small depth (so
|
||||
//! the producer stays a couple of batches ahead without unbounded
|
||||
//! memory growth).
|
||||
//!
|
||||
//! When the channel sender drops (either because all extents were
|
||||
//! served or because the [`Halt`] token cancelled), the consumer
|
||||
//! observes `RecvError` on the next `read_sectors` and treats it as
|
||||
//! end-of-stream. Errors from the underlying reader are forwarded
|
||||
//! verbatim through the channel.
|
||||
//!
|
||||
//! ## Read API
|
||||
//!
|
||||
//! `read_sectors` ignores its `lba`/`count` arguments — the producer
|
||||
//! has already chosen what to read, in the order the extents dictate.
|
||||
//! This is sound for the mux read path, which always walks extents
|
||||
//! sequentially and never seeks. For random-access callers (sweep
|
||||
//! patch retries) this wrapper is the wrong tool — they should keep
|
||||
//! reading the underlying source directly.
|
||||
|
||||
use crate::error::Result;
|
||||
use crate::event::{Event, EventKind};
|
||||
use crate::halt::Halt;
|
||||
use crate::sector::SectorSource;
|
||||
use crossbeam_channel::{Receiver, Sender, bounded};
|
||||
use std::thread::JoinHandle;
|
||||
|
||||
/// Producer-thread event callback. Fires `BytesRead` after every
|
||||
/// successful batch read so the consumer side can update a UI
|
||||
/// progress indicator without polling. `Send + 'static` because the
|
||||
/// callback runs on the producer thread.
|
||||
pub type EventFn = Box<dyn Fn(Event) + Send + 'static>;
|
||||
|
||||
const PREFETCH_CHANNEL_DEPTH: usize = 2;
|
||||
|
||||
/// Smallest sector source the producer will issue per read. AACS
|
||||
/// alignment requires multiples of 3 sectors so a unit doesn't span
|
||||
/// two reads.
|
||||
const SECTOR_ALIGNMENT: u16 = 3;
|
||||
|
||||
/// Item flowing through the prefetch forward channel.
|
||||
pub type Batch = std::result::Result<Vec<u8>, std::io::Error>;
|
||||
|
||||
/// Producer-thread-backed [`SectorSource`] decorator. Construct it
|
||||
/// with the real reader, the extent list to walk, and the batch
|
||||
/// size; the wrapper spawns the producer immediately and starts
|
||||
/// filling the channel.
|
||||
pub struct PrefetchedSectorSource {
|
||||
rx: Receiver<Batch>,
|
||||
/// Recycle channel — consumer returns drained buffers here; the
|
||||
/// producer re-fills them in place. Lets the producer/consumer
|
||||
/// reuse a fixed pool of `PREFETCH_CHANNEL_DEPTH+1` buffers
|
||||
/// instead of `Vec::new()`-ing one per batch (musl mallocng
|
||||
/// cross-thread alloc/free was the dominant cost in the demux
|
||||
/// thread before this).
|
||||
recycle_tx: Sender<Vec<u8>>,
|
||||
/// Joined on drop so producer cleanup runs deterministically.
|
||||
producer: Option<JoinHandle<()>>,
|
||||
/// Cumulative bytes drained by `read_sectors` calls. Exposed via
|
||||
/// [`capacity_sectors`] indirectly: the consumer-side state needs
|
||||
/// this to advance its position bookkeeping in lockstep with what
|
||||
/// the producer fed.
|
||||
total_sectors: u32,
|
||||
}
|
||||
|
||||
impl PrefetchedSectorSource {
|
||||
/// Spawn the producer thread. `reader` must already be the fully
|
||||
/// composed read+decrypt stack (e.g.
|
||||
/// [`DecryptingSectorSource`](crate::sector::DecryptingSectorSource))
|
||||
/// — every byte the producer emits is what the consumer's demux
|
||||
/// will feed to its codec parsers.
|
||||
pub fn new<S>(
|
||||
reader: S,
|
||||
extents: Vec<crate::disc::Extent>,
|
||||
batch_sectors: u16,
|
||||
halt: Option<Halt>,
|
||||
) -> Self
|
||||
where
|
||||
S: SectorSource + Send + 'static,
|
||||
{
|
||||
Self::new_with_events(reader, extents, batch_sectors, halt, None)
|
||||
}
|
||||
|
||||
/// Same as [`new`] but with a callback fired from the producer
|
||||
/// thread after each successful batch — used by autorip's mux
|
||||
/// path to surface `BytesRead` progress to the UI without the
|
||||
/// consumer thread having to poll.
|
||||
pub fn new_with_events<S>(
|
||||
mut reader: S,
|
||||
extents: Vec<crate::disc::Extent>,
|
||||
batch_sectors: u16,
|
||||
halt: Option<Halt>,
|
||||
event_fn: Option<EventFn>,
|
||||
) -> Self
|
||||
where
|
||||
S: SectorSource + Send + 'static,
|
||||
{
|
||||
let total_sectors: u32 = extents.iter().map(|e| e.sector_count).sum();
|
||||
let bytes_total_extents: u64 = extents.iter().map(|e| e.sector_count as u64 * 2048).sum();
|
||||
let (tx, rx) = bounded::<Batch>(PREFETCH_CHANNEL_DEPTH);
|
||||
let (recycle_tx, recycle_rx) = bounded::<Vec<u8>>(PREFETCH_CHANNEL_DEPTH + 1);
|
||||
let batch_bytes = batch_sectors as usize * 2048;
|
||||
|
||||
// Seed the recycle pool so the producer always has a buffer
|
||||
// to fill on the first iteration. Without these, the first
|
||||
// `recycle_rx.recv()` would block forever (no consumer has
|
||||
// returned a buffer yet).
|
||||
for _ in 0..(PREFETCH_CHANNEL_DEPTH + 1) {
|
||||
let _ = recycle_tx.send(vec![0u8; batch_bytes]);
|
||||
}
|
||||
|
||||
let producer = std::thread::Builder::new()
|
||||
.name("freemkv-prefetch".into())
|
||||
.spawn(move || {
|
||||
let mut ext_idx = 0usize;
|
||||
let mut offset: u32 = 0;
|
||||
let mut bytes_read_total: u64 = 0;
|
||||
while ext_idx < extents.len() {
|
||||
if halt.as_ref().map(|h| h.is_cancelled()).unwrap_or(false) {
|
||||
return;
|
||||
}
|
||||
let extent = &extents[ext_idx];
|
||||
let remaining = extent.sector_count.saturating_sub(offset);
|
||||
if remaining == 0 {
|
||||
ext_idx += 1;
|
||||
offset = 0;
|
||||
continue;
|
||||
}
|
||||
let mut sectors = remaining.min(batch_sectors as u32) as u16;
|
||||
if sectors >= SECTOR_ALIGNMENT {
|
||||
sectors -= sectors % SECTOR_ALIGNMENT;
|
||||
}
|
||||
let bytes = sectors as usize * 2048;
|
||||
let mut buf = match recycle_rx.recv() {
|
||||
Ok(b) => b,
|
||||
Err(_) => return, // consumer dropped both channels
|
||||
};
|
||||
if buf.len() < bytes {
|
||||
buf.resize(bytes, 0);
|
||||
} else {
|
||||
// Re-expose the full extent; previous truncate
|
||||
// shrank the visible len without freeing pages.
|
||||
// SAFETY: capacity is at least `bytes` after
|
||||
// construction with `vec![0u8; batch_bytes]`.
|
||||
unsafe { buf.set_len(bytes) };
|
||||
}
|
||||
let lba = extent.start_lba + offset;
|
||||
match reader.read_sectors(lba, sectors, &mut buf[..bytes], false) {
|
||||
Ok(n) => {
|
||||
buf.truncate(n);
|
||||
bytes_read_total = bytes_read_total.saturating_add(n as u64);
|
||||
if let Some(ref f) = event_fn {
|
||||
f(Event {
|
||||
kind: EventKind::BytesRead {
|
||||
bytes: bytes_read_total,
|
||||
total: bytes_total_extents,
|
||||
},
|
||||
});
|
||||
}
|
||||
if tx.send(Ok(buf)).is_err() {
|
||||
return; // consumer dropped
|
||||
}
|
||||
offset += sectors as u32;
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = tx.send(Err(e.into()));
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
// Drop tx implicitly — consumer sees RecvError → EOF.
|
||||
})
|
||||
.expect("freemkv-prefetch producer spawn failed");
|
||||
|
||||
Self {
|
||||
rx,
|
||||
recycle_tx,
|
||||
producer: Some(producer),
|
||||
total_sectors,
|
||||
}
|
||||
}
|
||||
|
||||
/// Peel off the receivers for zero-copy pipeline mode. The
|
||||
/// caller (typically [`super::super::mux::demux_thread::DemuxThread`])
|
||||
/// pulls buffers from `rx`, consumes them, and pushes the empty
|
||||
/// `Vec<u8>` back through `recycle_tx` so the producer can
|
||||
/// re-fill it. The producer-thread `JoinHandle` stays with the
|
||||
/// returned `PrefetchedSectorSource` shell; drop that to join.
|
||||
///
|
||||
/// Returns `(forward_rx, recycle_tx, shell)`. The shell only
|
||||
/// holds the join handle and total_sectors for `capacity_sectors`
|
||||
/// queries; its `SectorSource` impl becomes invalid after this
|
||||
/// call (data has been moved out).
|
||||
pub fn into_channels(self) -> (Receiver<Batch>, Sender<Vec<u8>>, PrefetchShell) {
|
||||
let total = self.total_sectors;
|
||||
// Drop the SectorSource side; transfer the producer join
|
||||
// handle to a shell that just waits on Drop.
|
||||
let mut me = self;
|
||||
let producer = me.producer.take();
|
||||
let rx = me.rx.clone();
|
||||
let recycle = me.recycle_tx.clone();
|
||||
std::mem::forget(me);
|
||||
(rx, recycle, PrefetchShell { producer, total })
|
||||
}
|
||||
}
|
||||
|
||||
/// Returned from [`PrefetchedSectorSource::into_channels`]. Owns the
|
||||
/// producer thread join handle so dropping the shell joins the
|
||||
/// producer, even though the channels have been peeled off.
|
||||
pub struct PrefetchShell {
|
||||
producer: Option<JoinHandle<()>>,
|
||||
#[allow(dead_code)]
|
||||
total: u32,
|
||||
}
|
||||
|
||||
impl Drop for PrefetchShell {
|
||||
fn drop(&mut self) {
|
||||
if let Some(h) = self.producer.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for PrefetchedSectorSource {
|
||||
fn drop(&mut self) {
|
||||
// Dropping the receiver closes the channel, which makes the
|
||||
// next producer `send` return Err and exits the loop. Joining
|
||||
// here gives us a deterministic shutdown — no detached thread
|
||||
// can outlive the source.
|
||||
if let Some(h) = self.producer.take() {
|
||||
let _ = h.join();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl SectorSource for PrefetchedSectorSource {
|
||||
fn capacity_sectors(&self) -> u32 {
|
||||
self.total_sectors
|
||||
}
|
||||
|
||||
fn read_sectors(
|
||||
&mut self,
|
||||
_lba: u32,
|
||||
_count: u16,
|
||||
buf: &mut [u8],
|
||||
_recovery: bool,
|
||||
) -> Result<usize> {
|
||||
// The producer has already decided the next batch. lba/count
|
||||
// are advisory; the consumer's fill_extents will advance its
|
||||
// own bookkeeping using the returned byte count, not the
|
||||
// requested count.
|
||||
match self.rx.recv() {
|
||||
Ok(Ok(filled)) => {
|
||||
let n = filled.len().min(buf.len());
|
||||
buf[..n].copy_from_slice(&filled[..n]);
|
||||
Ok(n)
|
||||
}
|
||||
Ok(Err(e)) => Err(crate::error::Error::IoError { source: e }),
|
||||
// Channel closed (producer finished or panicked).
|
||||
Err(_) => Ok(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -301,6 +301,7 @@ fn m2ts_stream_write_read() {
|
||||
pts: i as i64 * 1_000_000,
|
||||
keyframe: i == 0,
|
||||
data: vec![i; 100],
|
||||
duration_ns: None,
|
||||
};
|
||||
PesStream::write(&mut stream, &frame).unwrap();
|
||||
}
|
||||
@@ -322,6 +323,7 @@ fn m2ts_pes_frame_roundtrip() {
|
||||
pts: 1_234_567_890,
|
||||
keyframe: true,
|
||||
data: vec![0xDE; 200],
|
||||
duration_ns: None,
|
||||
};
|
||||
|
||||
let mut buf = Vec::new();
|
||||
@@ -589,6 +591,7 @@ fn mkvstream_write_finish() {
|
||||
pts: i as i64 * 1_000_000,
|
||||
keyframe: i == 0,
|
||||
data: vec![i; 100],
|
||||
duration_ns: None,
|
||||
};
|
||||
PesStream::write(&mut stream, &frame).unwrap();
|
||||
}
|
||||
@@ -650,6 +653,7 @@ fn mkvstream_roundtrip_bdts() {
|
||||
pts: i as i64 * 1_000_000,
|
||||
keyframe: true,
|
||||
data: vec![i; 100],
|
||||
duration_ns: None,
|
||||
};
|
||||
PesStream::write(&mut stream, &frame).unwrap();
|
||||
}
|
||||
@@ -842,6 +846,7 @@ fn mkvstream_e2e_h264_produces_valid_mkv() {
|
||||
pts: 1_000_000_000, // 1 second in ns
|
||||
keyframe: true,
|
||||
data: es_data,
|
||||
duration_ns: None,
|
||||
};
|
||||
PesStream::write(&mut stream2, &frame1).unwrap();
|
||||
|
||||
@@ -851,6 +856,7 @@ fn mkvstream_e2e_h264_produces_valid_mkv() {
|
||||
pts: 1_041_700_000, // ~1 frame later in ns
|
||||
keyframe: false,
|
||||
data: es_data2,
|
||||
duration_ns: None,
|
||||
};
|
||||
PesStream::write(&mut stream2, &frame2).unwrap();
|
||||
PesStream::finish(&mut stream2).unwrap();
|
||||
|
||||
Reference in New Issue
Block a user