From e2c84bfba71b28fbf14da908cc08d69c4b71c0b6 Mon Sep 17 00:00:00 2001 From: Radostin Stoyanov Date: Sat, 11 Jul 2026 12:12:40 +0100 Subject: [PATCH] mem/pagemap: Decode LZ4 pages before PIE The PIE restorer cannot link against LZ4. Decode mappings containing LZ4 payloads through CRIU's premap path, while retaining delayed PIE restore for raw and zero ranges when a worker pool cannot help. Classify image ranges as LZ4, raw, or zero. Decode LZ4 blocks while CRIU still runs in its normal address space, read aligned raw ranges directly into their final VMAs, and clear zero ranges in place. Omit compression metadata for entries containing only raw blocks so they retain the ordinary restore path. Share one encoded-read context across each incremental parent chain and reuse its worker pool for bounded batches, COW reads, streamed images, and page-server chunks. Use the same jobs to fill zero runs of at least 1 MiB when two blocks and two CPUs are available. Zero-only ranges that cannot benefit retain the PIE path, while zero blocks in an existing encoded batch remain inline without effective parallel capacity. Hold input buffers only while encoded work is active, and join every worker before child restore or the PIE handoff. Account calling restore threads and workers against a shared CPU budget. A requested width of zero selects automatic concurrency, one keeps LZ4 decoding serial and disables zero-fill workers, and larger values cap aggregate worker concurrency. Bound the active width by CPU affinity and useful work. Release encoded-work leases before ordinary reads so idle contexts cannot block ready workers. Local payload-heavy restores otherwise alternate between reading an encoded batch and decoding it. This leaves either storage or worker CPUs idle. When a second encoded-input lease is immediately available, dispatch the current batch to the pool and use the already-accounted calling thread to read the next adjacent payload. Keep synchronous reads for serial restore, streams, direct I/O, small batches, and contended leases. Select direct AIO from the ranges actually delayed for PIE. Batch inherited COW comparisons and streamed payload reads, and align page-server chunks to compression regions. Validate metadata and offset arithmetic before allocation or I/O, roll back partially enqueued work, and propagate worker, I/O, and deferred deduplication failures before remapping. On a 256 MiB pseudorandom workload with the worker limit set to eight and three measured runs, median restore time was 114.5 ms uncompressed, 113.4 ms with LZ4 per page, and 116.5 ms with 64 KiB regions. The previous compressed restore path took about 227 ms in either mode. With automatic concurrency and 256 KiB regions, parallel zero filling reduced median restore time over five runs from 103.1 ms to 31.7 ms for zero-filled memory and from 89.3 ms to 60.0 ms for mixed memory. On 512 MiB TEXT and ELF workloads with automatic concurrency and 256 KiB regions, five-run median restore time fell from 117.0 ms to 110.7 ms and from 115.4 ms to 109.5 ms, respectively. Assisted-by: Codex:GPT-5 Assisted-by: Claude:claude-fable-5 Signed-off-by: Radostin Stoyanov --- criu/compression.c | 822 +++++++++++++++ criu/config.c | 1 + criu/cr-restore.c | 16 + criu/include/compression.h | 126 +++ criu/include/page-pipe.h | 8 +- criu/include/page-xfer.h | 2 + criu/include/pagemap.h | 40 + criu/include/restorer.h | 16 + criu/include/rst_info.h | 2 + criu/mem.c | 211 +++- criu/page-pipe.c | 25 +- criu/page-xfer.c | 273 ++++- criu/pagemap.c | 1987 ++++++++++++++++++++++++++++-------- criu/pie/restorer.c | 353 +++++-- 14 files changed, 3306 insertions(+), 576 deletions(-) diff --git a/criu/compression.c b/criu/compression.c index d66264bc4..27e67eb61 100644 --- a/criu/compression.c +++ b/criu/compression.c @@ -1,16 +1,335 @@ +#include #include +#include +#include +#include #include #include +#include #include #include "page.h" #include "log.h" +#include "common/bug.h" #include "compression.h" +#include "common/xmalloc.h" #undef LOG_PREFIX #define LOG_PREFIX "compression: " +#define PARALLEL_DECOMPRESS_MIN_BYTES_PER_THREAD (512UL << 10) +#define PARALLEL_DECOMPRESS_STACK_SIZE (256UL << 10) +#define PARALLEL_DECOMPRESS_MAX_JOB_CHUNK 8 +#define PARALLEL_DECOMPRESS_MAX_BATCHES 2 + +/* + * asyncd can restore several shmem or memfd objects concurrently. Each of + * those requests may enter decompress_jobs_parallel_pool(), so limiting one + * invocation to the available CPUs is not sufficient: independent calls can + * otherwise multiply the worker count. Account for the calling thread as + * well as workers. Restore processes use one budget in shared restore memory; + * callers outside restore fall back to a process-wide budget. + */ +static pthread_once_t decompress_budget_once = PTHREAD_ONCE_INIT; +static pthread_mutex_t decompress_budget_lock = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t decompress_budget_cond = PTHREAD_COND_INITIALIZER; +static unsigned int decompress_budget_available = 1; +static unsigned int decompress_budget_capacity_cached = 1; +static pthread_mutex_t decompress_batch_lock = PTHREAD_MUTEX_INITIALIZER; +static pthread_cond_t decompress_batch_cond = PTHREAD_COND_INITIALIZER; +static unsigned int decompress_batches_available = PARALLEL_DECOMPRESS_MAX_BATCHES; +static struct decompression_shared_budget *decompress_shared_budget; + +static unsigned int decompression_available_cpus(void) +{ + cpu_set_t *affinity; + size_t affinity_size; + size_t nr_cpus = CPU_SETSIZE; + long available_cpus; + int affinity_errno; + + available_cpus = sysconf(_SC_NPROCESSORS_CONF); + if (available_cpus > (long)nr_cpus) + nr_cpus = (size_t)available_cpus; + + for (;;) { + affinity = CPU_ALLOC(nr_cpus); + if (!affinity) { + pr_warn("Unable to allocate a CPU affinity mask, using serial decompression\n"); + return 1; + } + affinity_size = CPU_ALLOC_SIZE(nr_cpus); + CPU_ZERO_S(affinity_size, affinity); + if (sched_getaffinity(0, affinity_size, affinity) == 0) { + available_cpus = CPU_COUNT_S(affinity_size, affinity); + CPU_FREE(affinity); + goto found; + } + affinity_errno = errno; + CPU_FREE(affinity); + if (affinity_errno != EINVAL) { + pr_warn("Unable to read the CPU affinity mask: %s; using serial decompression\n", + strerror(affinity_errno)); + return 1; + } + if (nr_cpus > UINT_MAX / 2) { + pr_warn("CPU affinity mask is too large, using serial decompression\n"); + return 1; + } + nr_cpus *= 2; + } +found: + if (available_cpus < 1) + return 1; + return (unsigned int)available_cpus; +} + +unsigned int decompression_thread_limit(unsigned int requested, unsigned int available_cpus) +{ + if (!available_cpus) + available_cpus = 1; + if (!requested) + return available_cpus; + return min(requested, available_cpus); +} + +static unsigned int decompression_cpu_limit(unsigned int requested) +{ + return decompression_thread_limit(requested, decompression_available_cpus()); +} + +/* + * Automatic (0) and serial (1) settings must not serialize independent asyncd + * jobs: both leave the restore-wide budget at the available-CPU capacity, and a + * serial setting still limits each worker-pool call to its caller. Explicit + * parallel settings cap all callers to the requested number of CPUs. + */ +static unsigned int decompression_budget_capacity(unsigned int configured_threads) +{ + if (configured_threads < 2) + return decompression_cpu_limit(0); + + return decompression_cpu_limit(configured_threads); +} + +static void decompress_budget_init(void) +{ + decompress_budget_capacity_cached = decompression_cpu_limit(0); + decompress_budget_available = decompress_budget_capacity_cached; +} + +void decompression_shared_budget_init(struct decompression_shared_budget *budget, unsigned int requested_threads) +{ + unsigned int threads = decompression_budget_capacity(requested_threads); + + /* Reduce an explicit request that exceeds the CPUs available to CRIU. */ + if (requested_threads > threads) + pr_warn("Reducing compressed-page worker concurrency from %u to %u (available CPUs)\n", + requested_threads, threads); + + futex_set(&budget->threads, threads); + budget->thread_capacity = threads; + /* + * Batch leases bound active encoded working sets independently of the CPU + * budget. Keep at most two 32 MiB input buffers so one local reader can + * overlap its next read with decoding without retaining unbounded memory. + */ + futex_set(&budget->batches, PARALLEL_DECOMPRESS_MAX_BATCHES); +} + +void decompression_use_shared_budget(struct decompression_shared_budget *budget) +{ + /* Borrowed from shared restore memory and valid for the restore lifetime. */ + decompress_shared_budget = budget; +} + +bool compressed_restore_has_parallel_capacity(unsigned int requested_threads) +{ + unsigned int capacity; + int err; + + if (requested_threads == 1) + return false; + if (decompress_shared_budget) + capacity = decompress_shared_budget->thread_capacity; + else { + err = pthread_once(&decompress_budget_once, decompress_budget_init); + if (err) + return false; + capacity = decompress_budget_capacity_cached; + } + + return decompression_thread_limit(requested_threads, capacity) > 1; +} + +static unsigned int shared_budget_acquire(futex_t *available, unsigned int requested) +{ + for (;;) { + unsigned int current = futex_get(available); + unsigned int granted; + unsigned int previous; + + if (!current) { + futex_wait_while_eq(available, 0); + continue; + } + granted = min(requested, current); + previous = atomic_cmpxchg(&available->raw, current, current - granted); + if (previous == current) + return granted; + } +} + +static bool shared_budget_try_acquire(futex_t *available) +{ + unsigned int current = futex_get(available); + + while (current) { + unsigned int previous; + + previous = atomic_cmpxchg(&available->raw, current, current - 1); + if (previous == current) + return true; + current = previous; + } + + return false; +} + +static void shared_budget_release(futex_t *available, unsigned int amount) +{ + atomic_add(amount, &available->raw); + futex_wake(available); +} + +static unsigned int decompress_budget_acquire(unsigned int requested) +{ + unsigned int granted; + int err; + + if (decompress_shared_budget) + return shared_budget_acquire(&decompress_shared_budget->threads, requested); + + err = pthread_once(&decompress_budget_once, decompress_budget_init); + if (err) { + pr_warn("Unable to initialize decompression CPU budget: %s\n", strerror(err)); + return 0; + } + + err = pthread_mutex_lock(&decompress_budget_lock); + if (err) { + pr_warn("Unable to lock decompression CPU budget: %s\n", strerror(err)); + return 0; + } + + while (!decompress_budget_available) { + err = pthread_cond_wait(&decompress_budget_cond, &decompress_budget_lock); + if (err) { + pr_warn("Unable to wait for decompression CPU budget: %s\n", strerror(err)); + pthread_mutex_unlock(&decompress_budget_lock); + return 0; + } + } + + granted = min(requested, decompress_budget_available); + decompress_budget_available -= granted; + pthread_mutex_unlock(&decompress_budget_lock); + return granted; +} + +static void decompress_budget_release(unsigned int nr_threads) +{ + int err; + + if (!nr_threads) + return; + if (decompress_shared_budget) { + shared_budget_release(&decompress_shared_budget->threads, nr_threads); + return; + } + + err = pthread_mutex_lock(&decompress_budget_lock); + if (err) { + pr_err("Unable to lock decompression CPU budget for release: %s\n", strerror(err)); + BUG(); + } + decompress_budget_available += nr_threads; + pthread_cond_broadcast(&decompress_budget_cond); + pthread_mutex_unlock(&decompress_budget_lock); +} + +static void decompress_budget_trim_reservation(unsigned int *threads_held, unsigned int threads_to_keep) +{ + BUG_ON(threads_to_keep > *threads_held); + decompress_budget_release(*threads_held - threads_to_keep); + *threads_held = threads_to_keep; +} + +void decompression_batch_acquire(void) +{ + int err; + + if (decompress_shared_budget) { + shared_budget_acquire(&decompress_shared_budget->batches, 1); + return; + } + err = pthread_mutex_lock(&decompress_batch_lock); + if (err) { + pr_err("Unable to lock decompression batch budget: %s\n", strerror(err)); + BUG(); + } + while (!decompress_batches_available) { + err = pthread_cond_wait(&decompress_batch_cond, &decompress_batch_lock); + if (err) { + pr_err("Unable to wait for decompression batch budget: %s\n", strerror(err)); + BUG(); + } + } + decompress_batches_available--; + pthread_mutex_unlock(&decompress_batch_lock); +} + +bool decompression_batch_try_acquire(void) +{ + bool acquired = false; + int err; + + if (decompress_shared_budget) + return shared_budget_try_acquire(&decompress_shared_budget->batches); + + err = pthread_mutex_lock(&decompress_batch_lock); + if (err) { + pr_warn("Unable to lock compressed-page batch budget: %s\n", strerror(err)); + return false; + } + if (decompress_batches_available) { + decompress_batches_available--; + acquired = true; + } + pthread_mutex_unlock(&decompress_batch_lock); + + return acquired; +} + +void decompression_batch_release(void) +{ + int err; + + if (decompress_shared_budget) { + shared_budget_release(&decompress_shared_budget->batches, 1); + return; + } + err = pthread_mutex_lock(&decompress_batch_lock); + if (err) { + pr_err("Unable to lock decompression batch budget for release: %s\n", strerror(err)); + BUG(); + } + decompress_batches_available++; + pthread_cond_signal(&decompress_batch_cond); + pthread_mutex_unlock(&decompress_batch_lock); +} + int compress_data(const char *input_data, size_t input_size, char *compressed_data, size_t output_size, int acceleration) @@ -165,3 +484,506 @@ int decompress_region(const char *src, int compressed_size, return 0; } + +/* Parallel decompression queue and reusable worker pool. */ +struct decompress_queue { + struct decompress_job *jobs; + size_t nr_jobs; + size_t job_chunk; + /* + * Concurrent scheduling and result accesses use compiler atomics. Keep all + * three fields under that API because next_job needs size_t, while CRIU's + * atomic_t is int-sized. + */ + size_t next_job; + int failed; + int failed_block; +}; + +struct decompress_worker { + pthread_t tid; + struct decompression_pool *pool; + unsigned int index; +}; + +/* + * The lock protects the batch scheduler fields below. decompression_pool_start_batch() + * publishes the caller-owned jobs as immutable batch metadata. Workers update + * only the queue's atomic scheduling and result fields. + * decompression_pool_finish_batch() waits for every selected worker before + * releasing those jobs. + */ +struct decompression_pool { + pthread_mutex_t lock; + pthread_cond_t work_ready; + pthread_cond_t work_done; + struct decompress_worker *workers; + struct decompress_queue queue; + unsigned int nr_workers; + unsigned int batch_workers; + unsigned int pending_batch_workers; + unsigned long batch_generation; + bool batch_active; + bool stop; +}; + +static int decompress_job_run(const struct decompress_job *job) +{ + size_t block_bytes; + + if (!job->pages || job->pages > MAX_REGION_PAGES) + return -1; + + block_bytes = (size_t)job->pages * PAGE_SIZE; + if (!job->compressed_size) { + memset(job->dst, 0, block_bytes); + return 0; + } + /* Raw blocks are handled inline by the pagemap reader. */ + if (job->compressed_size >= block_bytes) + return -1; + + return decompress_data_nolog(job->src, job->compressed_size, block_bytes, job->dst); +} + +static void decompress_queue_init(struct decompress_queue *queue, struct decompress_job *jobs, size_t nr_jobs, + unsigned int nr_threads) +{ + queue->jobs = jobs; + queue->nr_jobs = nr_jobs; + queue->job_chunk = nr_jobs / nr_threads / 4; + if (!queue->job_chunk) + queue->job_chunk = 1; + if (queue->job_chunk > PARALLEL_DECOMPRESS_MAX_JOB_CHUNK) + queue->job_chunk = PARALLEL_DECOMPRESS_MAX_JOB_CHUNK; + queue->next_job = 0; + queue->failed = 0; + queue->failed_block = INT_MAX; +} + +static bool decompress_queue_failed(const struct decompress_queue *queue) +{ + return __atomic_load_n(&queue->failed, __ATOMIC_ACQUIRE); +} + +static bool decompress_queue_claim_chunk(struct decompress_queue *queue, size_t *first, size_t *end) +{ + if (decompress_queue_failed(queue)) + return false; + + /* This relaxed counter only assigns disjoint chunks; it publishes no data. */ + *first = __atomic_fetch_add(&queue->next_job, queue->job_chunk, __ATOMIC_RELAXED); + if (*first >= queue->nr_jobs) + return false; + + *end = *first + queue->job_chunk; + if (*end > queue->nr_jobs) + *end = queue->nr_jobs; + return true; +} + +static void decompress_queue_record_failure(struct decompress_queue *queue, int block_index) +{ + int earliest_block = __atomic_load_n(&queue->failed_block, __ATOMIC_RELAXED); + + /* Keep the lowest failing block index when several workers fail. */ + while (block_index < earliest_block) { + if (__atomic_compare_exchange_n(&queue->failed_block, &earliest_block, block_index, false, + __ATOMIC_RELAXED, __ATOMIC_RELAXED)) + break; + } + + /* Publish cancellation only after failed_block has been updated. */ + __atomic_store_n(&queue->failed, 1, __ATOMIC_RELEASE); +} + +static void decompress_queue_run(struct decompress_queue *queue) +{ + size_t first, end; + + while (decompress_queue_claim_chunk(queue, &first, &end)) { + size_t i; + + for (i = first; i < end; i++) { + struct decompress_job *job = &queue->jobs[i]; + + if (!decompress_job_run(job)) + continue; + + decompress_queue_record_failure(queue, job->block_index); + break; + } + } +} + +static int decompress_jobs_serial(struct decompress_job *jobs, size_t nr_jobs) +{ + struct decompress_queue queue; + + decompress_queue_init(&queue, jobs, nr_jobs, 1); + decompress_queue_run(&queue); + if (!decompress_queue_failed(&queue)) + return 0; + + pr_err("Decompression failed at block %d\n", queue.failed_block); + return -1; +} + +static unsigned int decompress_batch_threads(size_t nr_jobs, size_t total_uncompressed, unsigned int requested_threads) +{ + size_t work_limit; + unsigned int nr_threads; + + if (total_uncompressed < PARALLEL_RESTORE_MIN_BATCH_BYTES) + return 1; + + nr_threads = decompression_cpu_limit(requested_threads); + work_limit = nr_jobs; + if (work_limit < nr_threads) + nr_threads = (unsigned int)work_limit; + work_limit = total_uncompressed / PARALLEL_DECOMPRESS_MIN_BYTES_PER_THREAD; + if (work_limit < nr_threads) + nr_threads = (unsigned int)work_limit; + if (nr_threads < 2) + return 1; + + return nr_threads; +} + +static void decompression_worker_signal_set(sigset_t *set) +{ + /* + * CRIU's process-global signal handling belongs to the restore thread. + * Worker threads keep asynchronous signals blocked for their lifetime. + * Synchronous faults must remain deliverable. + */ + sigfillset(set); + sigdelset(set, SIGABRT); + sigdelset(set, SIGBUS); + sigdelset(set, SIGFPE); + sigdelset(set, SIGILL); + sigdelset(set, SIGSEGV); + sigdelset(set, SIGSYS); + sigdelset(set, SIGTRAP); +} + +static void *decompression_pool_worker(void *arg) +{ + struct decompress_worker *worker = arg; + struct decompression_pool *pool = worker->pool; + unsigned long seen_generation = 0; + int err; + + err = pthread_mutex_lock(&pool->lock); + if (err) { + pr_err("Unable to lock decompression worker pool: %s\n", strerror(err)); + BUG(); + } + + for (;;) { + struct decompress_queue *queue; + + /* Wait for a new batch that selected this worker. */ + while (!pool->stop) { + if (seen_generation != pool->batch_generation && worker->index < pool->batch_workers) + break; + err = pthread_cond_wait(&pool->work_ready, &pool->lock); + if (err) { + pr_err("Unable to wait for decompression work: %s\n", strerror(err)); + BUG(); + } + } + if (pool->stop) + break; + + seen_generation = pool->batch_generation; + queue = &pool->queue; + pthread_mutex_unlock(&pool->lock); + + decompress_queue_run(queue); + + err = pthread_mutex_lock(&pool->lock); + if (err) { + pr_err("Unable to lock completed decompression work: %s\n", strerror(err)); + BUG(); + } + BUG_ON(!pool->pending_batch_workers); + pool->pending_batch_workers--; + if (!pool->pending_batch_workers) + pthread_cond_signal(&pool->work_done); + } + + pthread_mutex_unlock(&pool->lock); + return NULL; +} + +static void decompression_pool_spawn_workers(struct decompression_pool *pool, unsigned int nr_workers, + pthread_attr_t *attrp) +{ + unsigned int attempt; + + /* Successful workers are packed; a failed attempt does not consume a slot. */ + for (attempt = 0; attempt < nr_workers; attempt++) { + struct decompress_worker *worker = &pool->workers[pool->nr_workers]; + int err; + + worker->pool = pool; + worker->index = pool->nr_workers; + err = pthread_create(&worker->tid, attrp, decompression_pool_worker, worker); + if (err) { + pr_warn("Unable to start decompression worker %u: %s\n", attempt, strerror(err)); + break; + } + pool->nr_workers++; + } +} + +static struct decompression_pool *decompression_pool_create(unsigned int nr_threads, bool *mask_restore_failed) +{ + struct decompression_pool *pool; + pthread_attr_t attr; + sigset_t set, old; + unsigned int nr_workers; + size_t stack_size = PARALLEL_DECOMPRESS_STACK_SIZE; + int err; + + *mask_restore_failed = false; + if (nr_threads <= 1) + return NULL; + nr_workers = nr_threads - 1; + + pool = xzalloc(sizeof(*pool)); + if (!pool) + return NULL; + pool->workers = xzalloc((size_t)nr_workers * sizeof(*pool->workers)); + if (!pool->workers) + goto free_pool; + + /* Initialize every synchronization object before starting workers. */ + err = pthread_mutex_init(&pool->lock, NULL); + if (err) { + pr_warn("Unable to initialize decompression pool lock: %s\n", strerror(err)); + goto free_workers; + } + err = pthread_cond_init(&pool->work_ready, NULL); + if (err) { + pr_warn("Unable to initialize decompression work condition: %s\n", strerror(err)); + goto destroy_lock; + } + err = pthread_cond_init(&pool->work_done, NULL); + if (err) { + pr_warn("Unable to initialize decompression completion condition: %s\n", strerror(err)); + goto destroy_ready; + } + + /* A bounded worker stack keeps large restore pools inexpensive. */ + err = pthread_attr_init(&attr); + if (err) { + pr_warn("Unable to initialize decompression worker attributes: %s\n", strerror(err)); + goto destroy_done; + } + if (stack_size < PTHREAD_STACK_MIN) + stack_size = PTHREAD_STACK_MIN; + err = pthread_attr_setstacksize(&attr, stack_size); + if (err) { + pr_warn("Unable to set decompression worker stack size: %s\n", strerror(err)); + goto destroy_attr; + } + + /* Workers inherit blocked asynchronous signals from their creator. */ + decompression_worker_signal_set(&set); + err = pthread_sigmask(SIG_BLOCK, &set, &old); + if (err) { + pr_warn("Unable to block signals for decompression workers: %s\n", strerror(err)); + goto destroy_attr; + } + + decompression_pool_spawn_workers(pool, nr_workers, &attr); + + /* Restore the calling thread's mask after all workers have started. */ + err = pthread_sigmask(SIG_SETMASK, &old, NULL); + if (err) { + pr_err("Unable to restore signal mask after starting decompression workers: %s\n", strerror(err)); + *mask_restore_failed = true; + } + +destroy_attr: + pthread_attr_destroy(&attr); + + if (pool->nr_workers) + return pool; + +destroy_done: + pthread_cond_destroy(&pool->work_done); +destroy_ready: + pthread_cond_destroy(&pool->work_ready); +destroy_lock: + pthread_mutex_destroy(&pool->lock); +free_workers: + xfree(pool->workers); +free_pool: + xfree(pool); + return NULL; +} + +void decompression_pool_destroy(struct decompression_pool *pool) +{ + unsigned int w; + int err; + + if (!pool) + return; + + err = pthread_mutex_lock(&pool->lock); + if (err) { + pr_err("Unable to lock decompression pool for shutdown: %s\n", strerror(err)); + BUG(); + } + BUG_ON(pool->batch_active || pool->pending_batch_workers); + pool->stop = true; + pthread_cond_broadcast(&pool->work_ready); + pthread_mutex_unlock(&pool->lock); + + for (w = 0; w < pool->nr_workers; w++) { + err = pthread_join(pool->workers[w].tid, NULL); + if (err) { + pr_err("Unable to join decompression worker %u: %s\n", w, strerror(err)); + BUG(); + } + } + + pthread_cond_destroy(&pool->work_done); + pthread_cond_destroy(&pool->work_ready); + pthread_mutex_destroy(&pool->lock); + xfree(pool->workers); + xfree(pool); +} + +/* Publish one immutable job array to the selected subset of pool workers. */ +static void decompression_pool_start_batch(struct decompression_pool *pool, struct decompress_job *jobs, + size_t nr_jobs, unsigned int active_threads) +{ + int err; + + err = pthread_mutex_lock(&pool->lock); + if (err) { + pr_err("Unable to lock decompression pool for dispatch: %s\n", strerror(err)); + BUG(); + } + BUG_ON(pool->batch_active || pool->pending_batch_workers); + decompress_queue_init(&pool->queue, jobs, nr_jobs, active_threads); + pool->batch_active = true; + pool->batch_workers = active_threads - 1; + pool->pending_batch_workers = pool->batch_workers; + pool->batch_generation++; + pthread_cond_broadcast(&pool->work_ready); + pthread_mutex_unlock(&pool->lock); +} + +/* Wait until workers have dropped every reference to the caller-owned jobs. */ +static bool decompression_pool_finish_batch(struct decompression_pool *pool, int *failed_block) +{ + bool failed; + int err; + + err = pthread_mutex_lock(&pool->lock); + if (err) { + pr_err("Unable to lock decompression pool for completion: %s\n", strerror(err)); + BUG(); + } + while (pool->pending_batch_workers) { + err = pthread_cond_wait(&pool->work_done, &pool->lock); + if (err) { + pr_err("Unable to wait for decompression completion: %s\n", strerror(err)); + BUG(); + } + } + + failed = decompress_queue_failed(&pool->queue); + *failed_block = pool->queue.failed_block; + pool->queue.jobs = NULL; + pool->queue.nr_jobs = 0; + pool->batch_active = false; + pool->batch_workers = 0; + pthread_mutex_unlock(&pool->lock); + + return failed; +} + +int decompress_jobs_parallel_pool_with_caller_work( + struct decompression_pool **poolp, struct decompress_job *jobs, + size_t nr_jobs, size_t total_uncompressed, + unsigned int requested_threads, + decompression_caller_work_fn caller_work, void *caller_work_arg) +{ + struct decompression_pool *pool; + unsigned int active_threads, useful_threads, threads_held; + bool mask_restore_failed = false; + int failed_block = INT_MAX, ret = 0; + + if (!nr_jobs) + return 0; + if (!jobs || !poolp) + return -1; + + /* Reserve restore-wide CPU slots after choosing this batch's useful width. */ + useful_threads = decompress_batch_threads(nr_jobs, total_uncompressed, requested_threads); + threads_held = decompress_budget_acquire(useful_threads); + if (!threads_held) + return decompress_jobs_serial(jobs, nr_jobs); + if (threads_held == 1) { + ret = decompress_jobs_serial(jobs, nr_jobs); + goto out_release_budget; + } + + pool = *poolp; + if (pool && threads_held > pool->nr_workers + 1) { + decompression_pool_destroy(pool); + pool = NULL; + *poolp = NULL; + } + if (!pool) { + pool = decompression_pool_create(threads_held, &mask_restore_failed); + *poolp = pool; + } + if (mask_restore_failed) { + decompression_pool_destroy(pool); + *poolp = NULL; + ret = -1; + goto out_release_budget; + } + if (!pool) { + pr_warn("Unable to start decompression workers, using serial decompression\n"); + decompress_budget_trim_reservation(&threads_held, 1); + ret = decompress_jobs_serial(jobs, nr_jobs); + goto out_release_budget; + } + + /* A partially created pool may need fewer slots than were reserved. */ + active_threads = min(threads_held, pool->nr_workers + 1); + decompress_budget_trim_reservation(&threads_held, active_threads); + + decompression_pool_start_batch(pool, jobs, nr_jobs, active_threads); + if (caller_work) + caller_work(caller_work_arg); + /* After optional caller work, consume any jobs the pool has not claimed. */ + decompress_queue_run(&pool->queue); + + if (decompression_pool_finish_batch(pool, &failed_block)) { + pr_err("Decompression failed at block %d\n", failed_block); + ret = -1; + } +out_release_budget: + decompress_budget_release(threads_held); + return ret; +} + +int decompress_jobs_parallel_pool(struct decompression_pool **poolp, + struct decompress_job *jobs, size_t nr_jobs, + size_t total_uncompressed, + unsigned int requested_threads) +{ + return decompress_jobs_parallel_pool_with_caller_work( + poolp, jobs, nr_jobs, total_uncompressed, requested_threads, + NULL, NULL); +} diff --git a/criu/config.c b/criu/config.c index affdd8b36..ea75701ad 100644 --- a/criu/config.c +++ b/criu/config.c @@ -433,6 +433,7 @@ void init_opts(void) opts.image_io_mode = IMAGE_IO_DEFAULT; opts.network_lock_method = NETWORK_LOCK_DEFAULT; opts.ghost_fiemap = FIEMAP_DEFAULT; + opts.decompress_threads = 1; } bool deprecated_ok(char *what) diff --git a/criu/cr-restore.c b/criu/cr-restore.c index 6f3666d8c..3abdbbbf8 100644 --- a/criu/cr-restore.c +++ b/criu/cr-restore.c @@ -50,6 +50,7 @@ #include "uffd.h" #include "namespaces.h" #include "asyncd.h" +#include "compression.h" #include "mem.h" #include "mount.h" #include "fsnotify.h" @@ -2395,6 +2396,9 @@ int prepare_task_entries(void) mutex_init(&task_entries->userns_sync_lock); mutex_init(&task_entries->cgroupd_sync_lock); mutex_init(&task_entries->last_pid_mutex); + decompression_shared_budget_init(&task_entries->decompression_budget, + opts.decompress_threads); + decompression_use_shared_budget(&task_entries->decompression_budget); return 0; } @@ -3402,6 +3406,18 @@ static int sigreturn_restore(pid_t pid, struct task_restore_args *task_args, uns RST_MEM_FIXUP_PPTR(task_args->helpers); RST_MEM_FIXUP_PPTR(task_args->zombies); RST_MEM_FIXUP_PPTR(task_args->vma_ios); + { + struct restore_vma_io *rio = task_args->vma_ios; + unsigned int n; + + for (n = 0; n < task_args->vma_ios_n; n++) { + if (rio->compressed_size) + RST_MEM_FIXUP_PPTR(rio->compressed_size); + if (rio->block_pages) + RST_MEM_FIXUP_PPTR(rio->block_pages); + rio = (struct restore_vma_io *)((char *)rio + RIO_SIZE(rio->nr_iovs)); + } + } RST_MEM_FIXUP_PPTR(task_args->inotify_fds); task_args->compatible_mode = core_is_compat(core); diff --git a/criu/include/compression.h b/criu/include/compression.h index 8e36d280d..7b322c847 100644 --- a/criu/include/compression.h +++ b/criu/include/compression.h @@ -5,6 +5,7 @@ #include #include +#include "common/lock.h" #include "page.h" /* @@ -27,6 +28,9 @@ enum compress_mode { #define MAX_REGION_PAGES (MAX_REGION_SIZE / PAGE_SIZE) #define DEFAULT_REGION_PAGES (DEFAULT_REGION_SIZE / PAGE_SIZE) +/* Minimum useful batch before starting compressed-page restore workers. */ +#define PARALLEL_RESTORE_MIN_BATCH_BYTES (1UL << 20) + /* LZ4 worst-case compressed size for one page: src + src/255 + 16 */ #define PAGE_COMPRESSED_SIZE_BOUND (PAGE_SIZE + (PAGE_SIZE / 255) + 16) @@ -87,6 +91,34 @@ static inline bool page_is_all_zero(const char *page) return true; } +/* + * One block that can be restored independently. Jobs passed together must + * write to disjoint destinations. A zero compressed size clears the block; + * other jobs hold LZ4 payloads. Raw-fallback blocks remain caller-owned. + */ +struct decompress_job { + const char *src; + char *dst; + uint32_t compressed_size; + uint16_t pages; + int block_index; +}; + +struct decompression_pool; +typedef void (*decompression_caller_work_fn)(void *arg); + +/* + * Restore-wide limits live in the shared restore mapping. All task restore + * processes and asyncd workers inherit the same object, bounding active + * compressed-page workers and large encoded working sets across siblings. + */ +struct decompression_shared_budget { + futex_t threads; + futex_t batches; + /* Immutable capacity used for restore-path selection. */ + unsigned int thread_capacity; +}; + #ifdef CONFIG_LZ4 int compress_data(const char *input_data, size_t input_size, @@ -116,6 +148,50 @@ int compress_region(const char *src, unsigned int n_pages, char *dst, int decompress_region(const char *src, int compressed_size, unsigned int n_pages, char *dst); +/* + * Reuse worker threads through @pool across related batches. A later wider + * batch may replace the pool. @requested_threads includes the caller; zero + * selects automatic concurrency and one keeps the call serial. The active + * width is bounded by available CPUs, useful batch work, and the shared CPU + * budget. Small batches run serially. Calls sharing a pool must be serialized. + * @pool must initially be NULL and must be destroyed before the caller forks + * or remaps its address space. + */ +int decompress_jobs_parallel_pool(struct decompression_pool **pool, + struct decompress_job *jobs, + size_t nr_jobs, + size_t total_uncompressed, + unsigned int requested_threads); +/* + * After dispatching a parallel batch, run @caller_work once in the calling + * thread before it helps the pool. Serial fallbacks skip the callback. The + * callback must record its own result and must not mutate @jobs, re-enter + * @pool, or acquire another worker-budget reservation. + */ +int decompress_jobs_parallel_pool_with_caller_work( + struct decompression_pool **pool, struct decompress_job *jobs, + size_t nr_jobs, size_t total_uncompressed, + unsigned int requested_threads, + decompression_caller_work_fn caller_work, void *caller_work_arg); +void decompression_pool_destroy(struct decompression_pool *pool); +/* + * Initialize the restore-wide CPU and encoded-working-set budgets. Automatic + * (0) and serial (1) per-call settings still permit independent restore + * requests to run on separate CPUs; values above one cap their aggregate + * worker width. + */ +void decompression_shared_budget_init(struct decompression_shared_budget *budget, + unsigned int requested_threads); +void decompression_use_shared_budget(struct decompression_shared_budget *budget); +void decompression_batch_acquire(void); +bool decompression_batch_try_acquire(void); +void decompression_batch_release(void); +bool compressed_restore_has_parallel_capacity(unsigned int requested_threads); + +/* Apply the requested auto/explicit setting to the detected CPU capacity. */ +unsigned int decompression_thread_limit(unsigned int requested, + unsigned int available_cpus); + #else /* !CONFIG_LZ4 */ static inline int compress_data(const char *in, size_t in_sz, char *out, @@ -142,6 +218,56 @@ static inline int decompress_region(const char *src, int comp_sz, return -1; } +static inline int decompress_jobs_parallel_pool( + struct decompression_pool **pool, struct decompress_job *jobs, + size_t nr_jobs, size_t total_uncompressed, + unsigned int requested_threads) +{ + return -1; +} + +static inline int decompress_jobs_parallel_pool_with_caller_work( + struct decompression_pool **pool, struct decompress_job *jobs, + size_t nr_jobs, size_t total_uncompressed, + unsigned int requested_threads, + decompression_caller_work_fn caller_work, void *caller_work_arg) +{ + return -1; +} + +static inline void decompression_pool_destroy(struct decompression_pool *pool) +{ +} + +static inline void decompression_shared_budget_init( + struct decompression_shared_budget *budget, unsigned int requested_threads) +{ +} + +static inline void decompression_use_shared_budget( + struct decompression_shared_budget *budget) +{ +} + +static inline void decompression_batch_acquire(void) +{ +} + +static inline bool decompression_batch_try_acquire(void) +{ + return false; +} + +static inline void decompression_batch_release(void) +{ +} + +static inline bool compressed_restore_has_parallel_capacity( + unsigned int requested_threads) +{ + return false; +} + #endif /* CONFIG_LZ4 */ #endif /* __CR_COMPRESSION_H__ */ diff --git a/criu/include/page-pipe.h b/criu/include/page-pipe.h index 34709732b..cdd67689b 100644 --- a/criu/include/page-pipe.h +++ b/criu/include/page-pipe.h @@ -96,6 +96,7 @@ struct page_pipe_buf { unsigned long pipe_off; /* where this buf is started in a pipe */ unsigned long pages_in; /* how many pages are there */ #define PPB_LAZY (1 << 0) +#define PPB_FORCE_RAW (1 << 1) unsigned int flags; struct iovec *iov; /* vaddr:len map */ struct list_head l; /* links into page_pipe->bufs */ @@ -105,10 +106,11 @@ struct page_pipe_buf { * Page pipe buffers with different flags cannot share the same pipe. * We track the last ppb that was used for each type separately in the * prev[] array in the struct page_pipe (below). - * Currently we have 2 types: the buffers that are always stored in - * the images and the buffers that are lazily migrated + * PPB_LAZY and PPB_FORCE_RAW form the type index. Keeping force-raw + * buffers separate prevents a pipe segment that may be LZ4-compressed + * from being transferred under the force-raw policy (or vice versa). */ -#define PP_PIPE_TYPES 2 +#define PP_PIPE_TYPES 4 #define PP_HOLE_PARENT (1 << 0) diff --git a/criu/include/page-xfer.h b/criu/include/page-xfer.h index 6174daf56..202bc3ea2 100644 --- a/criu/include/page-xfer.h +++ b/criu/include/page-xfer.h @@ -31,6 +31,8 @@ struct page_xfer { */ unsigned long offset; bool transfer_lazy; + /* Store the current page-pipe segment as compressed-format raw/zero blocks. */ + bool force_raw; /* * Local non-streaming dump only: the pages image fd was switched to diff --git a/criu/include/pagemap.h b/criu/include/pagemap.h index 4154ee90b..23e6bc9f4 100644 --- a/criu/include/pagemap.h +++ b/criu/include/pagemap.h @@ -45,6 +45,8 @@ * All this is implemented in read_pagemap_page. */ +struct encoded_read_ctx; + struct page_read { /* reads page from current pagemap */ int (*read_pages)(struct page_read *, unsigned long vaddr, unsigned long nr, void *, unsigned flags); @@ -128,6 +130,15 @@ struct page_read { unsigned long cached_region_vaddr; size_t cached_region_size; + /* + * Bounded encoded buffers and workers. All readers in an incremental + * parent chain use the context owned by encoded_read_owner, so a sync + * initiated by a parent cannot acquire a second batch lease and deadlock + * against its child. Only the owner releases the context at close. + */ + struct encoded_read_ctx *encoded_read_ctx; + struct page_read *encoded_read_owner; + /* Record consequent neighbour iov-ecs to punch together */ struct iovec bunch; @@ -169,6 +180,35 @@ struct task_restore_args; int pagemap_enqueue_iovec(struct page_read *pr, void *buf, unsigned long len, struct list_head *to); int pagemap_render_iovec(struct list_head *from, struct task_restore_args *ta); +/* + * Return true when every file-backed queued read can be submitted with + * O_DIRECT. Encoded storage is rejected; raw storage and all destination + * extents must be page-aligned. Zero entries have no file extent, and a + * zero-only list returns false because it needs no AIO setup. + */ +bool pagemap_iovec_is_direct_compatible(const struct list_head *from); + +/* + * Return whether any page-image block overlapping [start, end) contains an + * actual LZ4 payload. Raw-fallback and zero blocks return false. Parent-image + * entries are followed without changing either reader's cursor. + * + * A local lookup is O(log(nr_pmes) + overlapping entries/blocks). Each + * inherited span repeats that lookup in the parent, so a fragmented parent + * chain additionally pays one logarithmic lookup per inherited span. A + * negative return indicates invalid input or an inconsistent parent chain. + */ +int page_read_range_has_lz4(struct page_read *pr, unsigned long start, + unsigned long end); + +/* Also return true when raw/zero runs exceed the bounded direct-PIE limit. */ +int page_read_range_needs_premap(struct page_read *pr, unsigned long start, + unsigned long end); + +/* Return whether the top image delegates any part of [start, end) to a parent. */ +int page_read_range_has_parent(struct page_read *pr, unsigned long start, + unsigned long end); + /* * Try to enable O_DIRECT on a pages-image fd and verify with one * aligned probe read. diff --git a/criu/include/restorer.h b/criu/include/restorer.h index 44a2f3c1e..2c3db3fc0 100644 --- a/criu/include/restorer.h +++ b/criu/include/restorer.h @@ -135,9 +135,25 @@ struct thread_restore_args { typedef long (*thread_restore_fcall_t)(struct thread_restore_args *args); +/* + * Internal representation of a pages-image range queued for PIE restore. + * + * PACKED_RAW and ZERO originate from entries which still carry compression + * metadata. They are separate from UNCOMPRESSED so that the restorer can + * bypass LZ4 without treating their packed image offsets as ordinary pages + * image offsets (in particular, --auto-dedup must not punch PACKED_RAW). + */ +enum restore_vma_io_storage { + VMA_IO_UNCOMPRESSED, + VMA_IO_ENCODED, + VMA_IO_PACKED_RAW, + VMA_IO_ZERO, +}; + struct restore_vma_io { int nr_iovs; loff_t off; + enum restore_vma_io_storage storage; uint32_t *compressed_size; uint64_t total_compressed_size; int n_compressed_size; diff --git a/criu/include/rst_info.h b/criu/include/rst_info.h index 691af2776..c2d7eb036 100644 --- a/criu/include/rst_info.h +++ b/criu/include/rst_info.h @@ -6,6 +6,7 @@ #include "common/list.h" #include "vma.h" #include "kerndat.h" +#include "compression.h" #include "images/mm.pb-c.h" #include "images/core.pb-c.h" @@ -17,6 +18,7 @@ struct task_entries { mutex_t userns_sync_lock; mutex_t cgroupd_sync_lock; mutex_t last_pid_mutex; + struct decompression_shared_budget decompression_budget; int asyncd_sk_id; }; diff --git a/criu/mem.c b/criu/mem.c index 29baf7599..0e398608d 100644 --- a/criu/mem.c +++ b/criu/mem.c @@ -229,9 +229,14 @@ static int generate_iovs(struct pstree_item *item, struct vma_area *vma, struct unsigned long pages[3] = {}; unsigned long vaddr; bool dump_all_pages; + bool self_contained; + bool force_raw; int ret = 0; + self_contained = (vma->e->flags & MAP_HUGETLB) || + (vma->e->status & VMA_EXT_PLUGIN); dump_all_pages = should_dump_entire_vma(vma->e); + force_raw = opts.compress_mode && self_contained; /* * In region-compression mode, force the first page of this VMA to @@ -262,6 +267,8 @@ static int generate_iovs(struct pstree_item *item, struct vma_area *vma, struct if (vma_entry_can_be_lazy(vma->e) && !is_stack(item, vaddr)) ppb_flags |= PPB_LAZY; + if (force_raw) + ppb_flags |= PPB_FORCE_RAW; /* * If we're doing incremental dump (parent images @@ -270,7 +277,15 @@ static int generate_iovs(struct pstree_item *item, struct vma_area *vma, struct * page. The latter would be checked in page-xfer. */ - if (has_parent && page_in_parent(page_info.softdirty)) { + /* + * Hugetlb and external-plugin VMAs cannot use the generic premap + * path, and delayed PIE I/O cannot redirect a PE_PARENT range to a + * different pages image. Keep these rare mappings self-contained at + * every image level. With compression enabled, force_raw additionally + * ensures PIE never has to decode LZ4. + */ + if (has_parent && !self_contained && + page_in_parent(page_info.softdirty)) { ret = page_pipe_add_hole(pp, vaddr, PP_HOLE_PARENT); st = 0; } else { @@ -1093,6 +1108,11 @@ static int premap_priv_vmas(struct pstree_item *t, struct vm_area_list *vmas, vo filemap_ctx_init(true); list_for_each_entry(vma, &vmas->h, list) { + bool exceptional; + int has_lz4 = 0; + int has_parent = 0; + int needs_premap = 0; + if (vma_area_is(vma, VMA_AREA_GUARD)) continue; @@ -1109,15 +1129,77 @@ static int premap_priv_vmas(struct pstree_item *t, struct vm_area_list *vmas, vo if (!vma_area_is_private(vma, kdat.task_size)) continue; + exceptional = (vma->e->flags & MAP_HUGETLB) || + (vma->e->status & VMA_EXT_PLUGIN); - if (vma->e->flags & MAP_HUGETLB) + /* + * PIE cannot link against liblz4. Premap ordinary private VMAs which + * contain an actual LZ4 block so their content is decoded by the + * normal page reader before switching to PIE. Also premap large zero + * runs through the compressed-page worker pool, and entries whose + * raw/zero run count exceeds the direct-PIE bound. Short runs then + * coalesce in-process, while long runs retain direct I/O. Uniform and + * lightly fragmented ranges keep the faster delayed path. + * + * Hugetlb and external-plugin VMAs cannot use the generic premap + * path. The dump side therefore guarantees that their compressed + * pagemap entries contain raw/zero blocks only. Reject an image + * which violates that invariant before the destructive restore. + */ + if (opts.compress_mode) { + if (exceptional) { + has_lz4 = page_read_range_has_lz4(pr, vma->e->start, vma->e->end); + if (has_lz4 < 0) { + ret = -1; + break; + } + } else { + needs_premap = page_read_range_needs_premap(pr, vma->e->start, vma->e->end); + if (needs_premap < 0) { + ret = -1; + break; + } + } + } + if (exceptional) { + has_parent = page_read_range_has_parent(pr, vma->e->start, vma->e->end); + if (has_parent < 0) { + ret = -1; + break; + } + if (has_parent) { + pr_err("Non-premapped VMA %#" PRIx64 "-%#" PRIx64 + " contains pages inherited from a parent image\n", + vma->e->start, vma->e->end); + ret = -1; + break; + } + } + + if (vma->e->flags & MAP_HUGETLB) { + if (has_lz4) { + pr_err("Hugetlb VMA %#" PRIx64 "-%#" PRIx64 + " contains an LZ4 block\n", + vma->e->start, vma->e->end); + ret = -1; + break; + } continue; + } /* VMA offset may change due to plugin so we cannot premap */ - if (vma->e->status & VMA_EXT_PLUGIN) + if (vma->e->status & VMA_EXT_PLUGIN) { + if (has_lz4) { + pr_err("External-plugin VMA %#" PRIx64 "-%#" PRIx64 + " contains an LZ4 block\n", + vma->e->start, vma->e->end); + ret = -1; + break; + } continue; + } - if (vma->pvma == NULL && pr->pieok && !vma_force_premap(vma, &vmas->h)) { + if (vma->pvma == NULL && pr->pieok && !needs_premap && !vma_force_premap(vma, &vmas->h)) { /* * VMA in question is not shared with anyone. We'll * restore it with its contents in restorer. @@ -1146,6 +1228,8 @@ static int premap_priv_vmas(struct pstree_item *t, struct vm_area_list *vmas, vo return ret; } +#define COW_READ_BATCH_PAGES (1UL << 8) + static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) { struct vma_area *vma; @@ -1162,18 +1246,11 @@ static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) unsigned int nr_lazy = 0; unsigned long va; void *buf = NULL; - int memerr; + bool page_read_closed = false; vma = list_first_entry(vmas, struct vma_area, list); rsti(t)->pages_img_id = pr->pages_img_id; - /* O_DIRECT may require the buffer to be aligned. */ - memerr = posix_memalign(&buf, PAGE_SIZE, PAGE_SIZE); - if (memerr) { - pr_err("Can't allocate COW buffer: %s\n", strerror(memerr)); - return -1; - } - /* * Read page contents. */ @@ -1198,7 +1275,8 @@ static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) continue; } - for (i = 0; i < nr_pages; i++) { + i = 0; + while (i < nr_pages) { void *p; /* @@ -1241,7 +1319,7 @@ static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) va += len; len >>= PAGE_SHIFT; nr_restored += len; - i += len - 1; + i += len; nr_enqueued++; continue; @@ -1254,24 +1332,64 @@ static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) off = (va - vma->e->start) / PAGE_SIZE; p = decode_pointer((off)*PAGE_SIZE + vma->premmaped_addr); - set_bit(off, vma->page_bitmap); if (vma_inherited(vma)) { - clear_bit(off, vma->pvma->page_bitmap); + unsigned long nr, j; + int memerr; - ret = pr->read_pages(pr, va, 1, buf, 0); - if (ret < 0) - goto err_read; + /* + * A page-at-a-time read turns compressed COW restore + * into one pread() and one LZ4 call per page. Read a + * bounded aligned batch, then retain the existing + * page-by-page sharing decision. + */ + nr = min_t(unsigned long, nr_pages - i, + (vma->e->end - va) / PAGE_SIZE); + nr = min(nr, COW_READ_BATCH_PAGES); + if (pr->pe->has_region_pages && pr->pe->region_pages && + nr < nr_pages - i) { + unsigned long aligned = + nr - nr % pr->pe->region_pages; - va += PAGE_SIZE; - nr_compared++; - - if (memcmp(p, buf, PAGE_SIZE) == 0) { - nr_shared++; /* the page is cowed */ - continue; + if (aligned) + nr = aligned; + } + if (!buf) { + memerr = posix_memalign(&buf, PAGE_SIZE, + COW_READ_BATCH_PAGES * PAGE_SIZE); + if (memerr) { + pr_err("Can't allocate COW buffer: %s\n", + strerror(memerr)); + ret = -1; + goto err_read; + } } - nr_restored++; - memcpy(p, buf, PAGE_SIZE); + ret = pr->read_pages(pr, va, nr, buf, PR_ASYNC); + if (ret < 0) + goto err_read; + if (pr->sync(pr)) { + ret = -1; + goto err_read; + } + + for (j = 0; j < nr; j++) { + void *src = (char *)buf + j * PAGE_SIZE; + void *dst = (char *)p + j * PAGE_SIZE; + + set_bit(off + j, vma->page_bitmap); + clear_bit(off + j, vma->pvma->page_bitmap); + nr_compared++; + if (memcmp(dst, src, PAGE_SIZE) == 0) { + nr_shared++; /* the page is cowed */ + continue; + } + + nr_restored++; + memcpy(dst, src, PAGE_SIZE); + } + + va += nr * PAGE_SIZE; + i += nr; } else { int nr; @@ -1286,13 +1404,14 @@ static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) nr = min_t(int, nr_pages - i, (vma->e->end - va) / PAGE_SIZE); + set_bit(off, vma->page_bitmap); ret = pr->read_pages(pr, va, nr, p, PR_ASYNC); if (ret < 0) goto err_read; va += nr * PAGE_SIZE; nr_restored += nr; - i += nr - 1; + i += nr; bitmap_set(vma->page_bitmap, off + 1, nr - 1); } @@ -1300,10 +1419,17 @@ static int restore_priv_vma_content(struct pstree_item *t, struct page_read *pr) } err_read: - if (pr->sync(pr)) - goto out; + { + int sync_ret = pr->sync(pr); + + pr->close(pr); + page_read_closed = true; + if (sync_ret) { + ret = -1; + goto out; + } + } - pr->close(pr); if (ret < 0) { exit_code = ret; goto out; @@ -1354,6 +1480,11 @@ err_read: err_addr: pr_err("Page entry address %lx outside of VMA %lx-%lx\n", va, (long)vma->e->start, (long)vma->e->end); out: + if (!page_read_closed) { + if (pr->sync(pr)) + exit_code = -1; + pr->close(pr); + } xfree(buf); return exit_code; } @@ -1529,6 +1660,7 @@ int open_vmas(struct pstree_item *t) static int prepare_vma_ios(struct pstree_item *t, struct task_restore_args *ta) { struct cr_img *pages; + int ret; /* * We optimize the case when rsti(t)->vma_io is empty. @@ -1554,15 +1686,28 @@ static int prepare_vma_ios(struct pstree_item *t, struct task_restore_args *ta) return -1; ta->vma_ios_fd = img_raw_fd(pages); - if (ta->vma_ios_fd >= 0 && opts.image_io_mode == IMAGE_IO_DIRECT) { + /* + * Select direct AIO from the actual delayed ranges, not the inventory-wide + * compression mode. LZ4 ranges were premapped above, while aligned raw + * fallbacks and zero ranges remain safe for the PIE fast path. + */ + if (ta->vma_ios_fd >= 0 && opts.image_io_mode == IMAGE_IO_DIRECT && + pagemap_iovec_is_direct_compatible(&rsti(t)->vma_io)) { int direct = probe_pages_o_direct(ta->vma_ios_fd); if (direct < 0) { close_image(pages); + ta->vma_ios_fd = -1; return -1; } ta->vma_ios_use_direct = (direct == 1); } - return pagemap_render_iovec(&rsti(t)->vma_io, ta); + + ret = pagemap_render_iovec(&rsti(t)->vma_io, ta); + if (ret) { + close_image(pages); + ta->vma_ios_fd = -1; + } + return ret; } int prepare_vmas(struct pstree_item *t, struct task_restore_args *ta) diff --git a/criu/page-pipe.c b/criu/page-pipe.c index 8fe16e318..a64bf4785 100644 --- a/criu/page-pipe.c +++ b/criu/page-pipe.c @@ -68,9 +68,20 @@ static inline int ppb_resize_pipe(struct page_pipe_buf *ppb) return 0; } +static inline unsigned int ppb_type(unsigned int ppb_flags) +{ + unsigned int type = 0; + + if (ppb_flags & PPB_LAZY && opts.lazy_pages) + type |= PPB_LAZY; + if (ppb_flags & PPB_FORCE_RAW) + type |= PPB_FORCE_RAW; + + return type; +} + static struct page_pipe_buf *pp_prev_ppb(struct page_pipe *pp, unsigned int ppb_flags) { - int type = 0; /* don't allow to reuse a pipe in the PP_CHUNK_MODE mode */ if (pp->flags & PP_CHUNK_MODE) @@ -79,20 +90,12 @@ static struct page_pipe_buf *pp_prev_ppb(struct page_pipe *pp, unsigned int ppb_ if (list_empty(&pp->bufs)) return NULL; - if (ppb_flags & PPB_LAZY && opts.lazy_pages) - type = 1; - - return pp->prev[type]; + return pp->prev[ppb_type(ppb_flags)]; } static void pp_update_prev_ppb(struct page_pipe *pp, struct page_pipe_buf *ppb, unsigned int ppb_flags) { - int type = 0; - - if (ppb_flags & PPB_LAZY && opts.lazy_pages) - type = 1; - - pp->prev[type] = ppb; + pp->prev[ppb_type(ppb_flags)] = ppb; } static struct page_pipe_buf *ppb_alloc(struct page_pipe *pp, unsigned int ppb_flags) diff --git a/criu/page-xfer.c b/criu/page-xfer.c index 9523f3eea..0f75b3212 100644 --- a/criu/page-xfer.c +++ b/criu/page-xfer.c @@ -8,6 +8,7 @@ #include #include #include +#include #undef LOG_PREFIX #define LOG_PREFIX "page-xfer: " @@ -264,6 +265,8 @@ static int write_pages_to_server_compressed(struct page_xfer *xfer, int p, unsig if (page_is_all_zero(buf)) { compressed_size = 0; + } else if (xfer->force_raw) { + compressed_size = PAGE_SIZE; } else { int r = compress_data(buf, PAGE_SIZE, compressed_buf, PAGE_COMPRESSED_SIZE_BOUND, acceleration); if (r < 0) @@ -540,6 +543,40 @@ static int read_pipe_full(int fd, void *buf, size_t count) return 0; } +static bool region_is_all_zero(const char *buf, unsigned int nr_pages) +{ + unsigned int i; + + for (i = 0; i < nr_pages; i++) + if (!page_is_all_zero(buf + (size_t)i * PAGE_SIZE)) + return false; + + return true; +} + +static bool pending_entry_is_all_raw(const struct page_xfer *xfer) +{ + unsigned int region_pages = xfer->pending_pe.region_pages; + size_t i; + + for (i = 0; i < xfer->pending_pe.total_blocks; i++) { + size_t block_bytes = PAGE_SIZE; + + if (region_pages) { + unsigned long pages_done = i * (unsigned long)region_pages; + unsigned long pages_left = xfer->pending_pe.nr_pages - pages_done; + unsigned int block_pages = pages_left < region_pages ? (unsigned int)pages_left : region_pages; + + block_bytes = (size_t)block_pages * PAGE_SIZE; + } + + if (xfer->pending_pe.compressed_size[i] != block_bytes) + return false; + } + + return true; +} + static int write_pages_loc_compressed(struct page_xfer *xfer, int p, unsigned long len) { unsigned long off = 0; @@ -578,6 +615,13 @@ static int write_pages_loc_compressed(struct page_xfer *xfer, int p, unsigned lo xfer->pending_pe.compressed_size[idx] = 0; continue; } + if (xfer->force_raw) { + xfer->pending_pe.compressed_size[idx] = PAGE_SIZE; + xfer->pending_pe.total_compressed_size += PAGE_SIZE; + if (write_compressed_payload(xfer, buf, PAGE_SIZE, true)) + return -1; + continue; + } cs = compress_data(buf, PAGE_SIZE, compressed_buf, PAGE_COMPRESSED_SIZE_BOUND, acceleration); if (cs < 0) @@ -608,8 +652,8 @@ static int write_pages_loc_compressed(struct page_xfer *xfer, int p, unsigned lo int rc = -1; src_buf = xmalloc(region_bytes_max); - dst_buf = xmalloc(cap); - if (!src_buf || !dst_buf) + dst_buf = xfer->force_raw ? NULL : xmalloc(cap); + if (!src_buf || (!xfer->force_raw && !dst_buf)) goto region_out; pages_done = (unsigned long)xfer->pending_pe.n_compressed * region_pages; @@ -618,6 +662,7 @@ static int write_pages_loc_compressed(struct page_xfer *xfer, int p, unsigned lo unsigned long pages_left = xfer->pending_pe.nr_pages - pages_done; unsigned int this_region = pages_left < region_pages ? pages_left : region_pages; size_t region_bytes = (size_t)this_region * PAGE_SIZE; + const char *payload = dst_buf; int cs; size_t idx; @@ -633,13 +678,22 @@ static int write_pages_loc_compressed(struct page_xfer *xfer, int p, unsigned lo idx = xfer->pending_pe.n_compressed++; - cs = compress_region(src_buf, this_region, dst_buf, cap, acceleration); - if (cs < 0) - goto region_out; + if (xfer->force_raw) { + if (region_is_all_zero(src_buf, this_region)) { + cs = 0; + } else { + cs = (int)region_bytes; + payload = src_buf; + } + } else { + cs = compress_region(src_buf, this_region, dst_buf, cap, acceleration); + if (cs < 0) + goto region_out; + } xfer->pending_pe.compressed_size[idx] = cs; xfer->pending_pe.total_compressed_size += cs; - if (cs > 0 && write_compressed_payload(xfer, dst_buf, cs, (size_t)cs == region_bytes)) + if (cs > 0 && write_compressed_payload(xfer, payload, cs, (size_t)cs == region_bytes)) goto region_out; } @@ -654,20 +708,28 @@ region_out: /* When all blocks are compressed, flush the pagemap entry */ if (xfer->pending_pe.n_compressed == xfer->pending_pe.total_blocks) { PagemapEntry pe = PAGEMAP_ENTRY__INIT; + bool all_raw = pending_entry_is_all_raw(xfer); pe.vaddr = xfer->pending_pe.vaddr; pe.nr_pages = xfer->pending_pe.nr_pages; pe.has_flags = true; pe.flags = xfer->pending_pe.flags; pe.has_nr_pages = true; - pe.compressed_size = xfer->pending_pe.compressed_size; - pe.n_compressed_size = xfer->pending_pe.total_blocks; - pe.has_total_compressed_size = true; - pe.total_compressed_size = xfer->pending_pe.total_compressed_size; + /* + * An entry for which every block fell back to raw has exactly the + * same contiguous payload as an ordinary page image. Omitting the + * compression metadata lets restore take the uncompressed fast path. + */ + if (!all_raw) { + pe.compressed_size = xfer->pending_pe.compressed_size; + pe.n_compressed_size = xfer->pending_pe.total_blocks; + pe.has_total_compressed_size = true; + pe.total_compressed_size = xfer->pending_pe.total_compressed_size; - if (region_pages > 0) { - pe.has_region_pages = true; - pe.region_pages = region_pages; + if (region_pages > 0) { + pe.has_region_pages = true; + pe.region_pages = region_pages; + } } if (pb_write_one(xfer->pmi, &pe, PB_PAGEMAP) < 0) @@ -876,6 +938,7 @@ out: xfer->pending_pe.compressed_size = NULL; xfer->pending_pe.payload_started = false; xfer->pages_image_offset = 0; + xfer->force_raw = false; if (compress) { xfer->write_pagemap = write_pagemap_loc_compressed; xfer->write_pages = write_pages_loc_compressed; @@ -899,6 +962,7 @@ int open_page_xfer(struct page_xfer *xfer, int fd_type, unsigned long img_id) xfer->offset = 0; xfer->transfer_lazy = true; + xfer->force_raw = false; if (opts.use_page_server) return open_page_server_xfer(xfer, fd_type, img_id); @@ -1337,6 +1401,7 @@ int page_xfer_predump_pages(int pid, struct page_xfer *xfer, struct page_pipe *p pr_debug("\t p %p - %p\n", iov.iov_base, iov.iov_base + iov.iov_len); flags = ppb_xfer_flags(xfer, ppb); + xfer->force_raw = ppb->flags & PPB_FORCE_RAW; if (xfer->write_pagemap(xfer, &iov, flags)) goto err; @@ -1385,6 +1450,7 @@ int page_xfer_dump_pages(struct page_xfer *xfer, struct page_pipe *pp) pr_debug("\tp %p - %p\n", iov.iov_base, iov.iov_base + iov.iov_len); flags = ppb_xfer_flags(xfer, ppb); + xfer->force_raw = ppb->flags & PPB_FORCE_RAW; if (xfer->write_pagemap(xfer, &iov, flags)) return -1; @@ -1601,6 +1667,7 @@ static int page_server_add_compressed(int sk, struct page_server_iov *pi, u32 fl unsigned long i; uint32_t *compressed_size; uint64_t total_compressed_size = 0; + bool all_raw = true; bool payload_started = false; if (validate_page_server_iov(pi, flags)) @@ -1649,6 +1716,8 @@ static int page_server_add_compressed(int sk, struct page_server_iov *pi, u32 fl compressed_size[i] = cs; total_compressed_size += cs; + if (cs != PAGE_SIZE) + all_raw = false; if (cs > 0) { char buf[PAGE_COMPRESSED_SIZE_BOUND]; @@ -1681,10 +1750,12 @@ static int page_server_add_compressed(int sk, struct page_server_iov *pi, u32 fl pe.has_flags = true; pe.flags = flags; pe.has_nr_pages = true; - pe.compressed_size = compressed_size; - pe.n_compressed_size = pi->nr_pages; - pe.has_total_compressed_size = true; - pe.total_compressed_size = total_compressed_size; + if (!all_raw) { + pe.compressed_size = compressed_size; + pe.n_compressed_size = pi->nr_pages; + pe.has_total_compressed_size = true; + pe.total_compressed_size = total_compressed_size; + } if (pb_write_one(lxfer->pmi, &pe, PB_PAGEMAP) < 0) { xfree(compressed_size); @@ -1969,11 +2040,153 @@ static int page_server_serve(int sk) return ret; } -static int fill_page_pipe(struct page_read *pr, struct page_pipe *pp) +/* Match the restore batch cap; individual reads are aligned below. */ +#define PAGE_SERVER_DECODE_BUFFER_SIZE (32UL << 20) + +static bool page_read_requires_buffered_copy(const struct page_read *pr) +{ + int i; + + /* splice(2) is not portable for O_DIRECT input descriptors. */ + if (pr->use_direct) + return true; + + for (i = 0; i < pr->nr_pmes; i++) + if (pagemap_present(pr->pmes[i]) && + (pr->pmes[i]->n_compressed_size || + pagemap_payload_aligned(pr->pmes[i]))) + return true; + + return false; +} + +static int splice_page_pipe(struct page_read *pr, struct page_pipe *pp) { struct page_pipe_buf *ppb; + int fd = img_raw_fd(pr->pi); + int i; + + if (fd < 0) { + pr_err("Failed getting pages image fd\n"); + return -1; + } + + list_for_each_entry(ppb, &pp->bufs, l) { + for (i = 0; i < ppb->nr_segs; i++) { + size_t left = ppb->iov[i].iov_len; + + while (left) { + ssize_t ret = splice(fd, NULL, ppb->p[1], NULL, left, SPLICE_F_MOVE); + + if (ret < 0) { + if (errno == EINTR) + continue; + pr_perror("Splice from pages image failed"); + return -1; + } + if (!ret) { + pr_err("Pages image ended while filling page pipe\n"); + return -1; + } + left -= ret; + } + } + } + + return 0; +} + +static int decode_page_pipe(struct page_read *pr, struct page_pipe *pp) +{ + const size_t buffer_size = max((size_t)PAGE_SIZE, (size_t)PAGE_SERVER_DECODE_BUFFER_SIZE); + const unsigned long buffer_pages = buffer_size / PAGE_SIZE; + struct page_pipe_buf *ppb; + void *buf = NULL; int i, ret; + ret = posix_memalign(&buf, PAGE_SIZE, buffer_size); + if (ret) { + pr_err("Failed to allocate page-pipe decode buffer: %s\n", strerror(ret)); + return -1; + } + + pr->reset(pr); + list_for_each_entry(ppb, &pp->bufs, l) { + for (i = 0; i < ppb->nr_segs; i++) { + struct iovec iov = ppb->iov[i]; + unsigned long vaddr = encode_pointer(iov.iov_base); + unsigned long nr_pages = iov.iov_len / PAGE_SIZE; + unsigned long pages_done = 0; + + if (iov.iov_len % PAGE_SIZE) { + pr_err("Unaligned page-pipe iovec length %zu\n", iov.iov_len); + goto err; + } + + while (pages_done < nr_pages) { + unsigned int region_pages; + unsigned long entry_pages; + unsigned long read_pages; + unsigned long page_vaddr = vaddr + pages_done * PAGE_SIZE; + + ret = pr->seek_pagemap(pr, page_vaddr); + if (ret <= 0 || !pr->pe || pagemap_in_parent(pr->pe)) { + pr_err("Missing local page at %#lx while filling page pipe\n", page_vaddr); + goto err; + } + + entry_pages = pr->pe->nr_pages - ((page_vaddr - pr->pe->vaddr) / PAGE_SIZE); + read_pages = entry_pages; + read_pages = min(read_pages, nr_pages - pages_done); + region_pages = pr->pe->has_region_pages ? pr->pe->region_pages : 0; + if (read_pages > buffer_pages || + (region_pages && read_pages == buffer_pages && + entry_pages > read_pages)) { + unsigned long batch_pages = buffer_pages; + + /* + * Region sizes are arbitrary page multiples, so a fixed + * 32 MiB buffer is not divisible by all of them. End a + * bounded chunk on a region boundary; otherwise the next + * chunk falls back to synchronous partial decompression. + */ + if (region_pages) + batch_pages -= batch_pages % region_pages; + if (!batch_pages) { + pr_err("Compression region %u exceeds page-server decode buffer\n", + region_pages); + goto err; + } + read_pages = batch_pages; + } + ret = pr->read_pages(pr, page_vaddr, read_pages, buf, PR_ASYNC); + if (ret < 0) + goto err; + if (pr->sync(pr)) + goto err; + + if (write_fd_full(ppb->p[1], buf, read_pages * PAGE_SIZE)) { + pr_err("Failed writing decoded pages to page pipe\n"); + goto err; + } + pages_done += read_pages; + } + } + } + + xfree(buf); + return 0; + +err: + xfree(buf); + return -1; +} + +static int fill_page_pipe(struct page_read *pr, struct page_pipe *pp) +{ + unsigned long i; + int ret; + pr->reset(pr); while (pr->advance(pr)) { @@ -1991,17 +2204,17 @@ static int fill_page_pipe(struct page_read *pr, struct page_pipe *pp) } } - list_for_each_entry(ppb, &pp->bufs, l) { - for (i = 0; i < ppb->nr_segs; i++) { - struct iovec iov = ppb->iov[i]; - - if (splice(img_raw_fd(pr->pi), NULL, ppb->p[1], NULL, iov.iov_len, SPLICE_F_MOVE) != - iov.iov_len) { - pr_perror("Splice failed"); - return -1; - } - } - } + /* + * Ordinary page images retain the original zero-copy send path. A page + * image with compression metadata needs the page reader to synthesize + * zero blocks, copy raw fallbacks, and decode LZ4 blocks. + */ + if (page_read_requires_buffered_copy(pr)) + ret = decode_page_pipe(pr, pp); + else + ret = splice_page_pipe(pr, pp); + if (ret) + return ret; debug_show_page_pipe(pp); @@ -2018,6 +2231,8 @@ static int page_pipe_from_pagemap(struct page_pipe **pp, int pid) pr_err("Failed to open page read for %d\n", pid); return -1; } + /* Building the serving pipe must not modify the source image on failure. */ + page_read_disable_dedup(&pr); while (pr.advance(&pr)) if (pagemap_present(pr.pe)) diff --git a/criu/pagemap.c b/criu/pagemap.c index a6a2d955c..81c65d89e 100644 --- a/criu/pagemap.c +++ b/criu/pagemap.c @@ -35,6 +35,14 @@ #define ASYNC_BATCH_MAX_PAGES (ASYNC_BATCH_MAX_BYTES / PAGE_SIZE) #define ASYNC_READAHEAD_MAX_GAP (1UL << 20) #define ASYNC_READAHEAD_MAX_BYTES (256UL << 20) +#define DIRECT_COMPRESSED_RUN_MAX 128 +/* + * Splitting a tiny raw/zero island out of an encoded batch trades a memcpy for + * another restore job (and, for raw data, another preadv()). Keep small + * islands with neighbouring LZ4 blocks, while always exposing raw/zero ranges + * when the whole request can be restored without LZ4. + */ +#define DIRECT_COMPRESSED_RUN_MIN_BYTES (64UL * 1024) #define OFF_MAX (sizeof(off_t) == sizeof(long long) ? LLONG_MAX : sizeof(off_t) == sizeof(int) ? INT_MAX : -999999) #define OFF_MIN (sizeof(off_t) == sizeof(long long) ? LLONG_MIN : sizeof(off_t) == sizeof(int) ? INT_MIN : -999999) @@ -47,6 +55,7 @@ struct page_read_iov { off_t end; /* exclusive end offset in the pages image */ struct iovec *to; /* destination iovs */ unsigned int nr; /* their number */ + enum restore_vma_io_storage storage; size_t n_compressed_size; /* Number of compressed blocks (pages or regions) */ uint32_t *compressed_size; /* Per-block compressed sizes */ @@ -80,6 +89,77 @@ struct page_read_iov { struct list_head l; }; +struct encoded_prefetch { + char *buffer; + size_t count; + size_t done; + off_t offset; + int fd; + int saved_errno; + bool complete; +}; + +/* + * Reusable storage for encoded reads. The top-level page reader owns one + * context for its whole parent chain. Buffers are reused during one active + * read and then released; decompression workers remain reusable across reads. + */ +struct encoded_read_ctx { + struct decompress_job *jobs; + size_t jobs_cap; + char *compressed; + size_t compressed_cap; + char *prefetch_buffer; + size_t prefetch_cap; + struct page_read_iov *prefetched_piov; + struct encoded_prefetch prefetch; + char *scratch; + size_t scratch_cap; + struct decompression_pool *pool; + bool batch_acquired; + bool prefetch_batch_acquired; +}; + +/* Reserve one restore-wide encoded working set for this active read. */ +static void encoded_read_ctx_begin_work(struct encoded_read_ctx *ctx) +{ + BUG_ON(ctx->batch_acquired); + decompression_batch_acquire(); + ctx->batch_acquired = true; +} + +/* + * Return the batch slot only after its process-local buffers are gone. Worker + * threads remain reusable, but idle page readers no longer pin a 32 MiB slot + * or retain memory outside the restore-wide bound. + */ +static void encoded_read_ctx_end_work(struct encoded_read_ctx *ctx) +{ + if (!ctx || !ctx->batch_acquired) + return; + + xfree(ctx->scratch); + ctx->scratch = NULL; + ctx->scratch_cap = 0; + xfree(ctx->compressed); + ctx->compressed = NULL; + ctx->compressed_cap = 0; + xfree(ctx->prefetch_buffer); + ctx->prefetch_buffer = NULL; + ctx->prefetch_cap = 0; + ctx->prefetched_piov = NULL; + memset(&ctx->prefetch, 0, sizeof(ctx->prefetch)); + xfree(ctx->jobs); + ctx->jobs = NULL; + ctx->jobs_cap = 0; + if (ctx->prefetch_batch_acquired) { + ctx->prefetch_batch_acquired = false; + decompression_batch_release(); + } + ctx->batch_acquired = false; + decompression_batch_release(); +} + static inline bool can_extend_bunch(struct iovec *bunch, unsigned long off, unsigned long len) { return /* The next region is the continuation of the existing */ @@ -436,21 +516,78 @@ static int read_local_page(struct page_read *pr, unsigned long vaddr, unsigned l return 0; } +static unsigned int pagemap_region_pages(const PagemapEntry *pe) +{ + if (pe && pe->has_region_pages && pe->region_pages) + return pe->region_pages; + return 0; +} + +static int compressed_block_layout(const PagemapEntry *pe, size_t index, + unsigned int *block_pages, size_t *block_bytes) +{ + unsigned int region_pages = pagemap_region_pages(pe); + uint64_t first_page; + uint64_t pages_left; + + if (!pe || index >= pe->n_compressed_size) { + pr_err("Compressed block index %zu is out of bounds\n", index); + return -1; + } + + if (!region_pages) { + *block_pages = 1; + *block_bytes = PAGE_SIZE; + return 0; + } + + first_page = (uint64_t)index * region_pages; + if (first_page >= pe->nr_pages) { + pr_err("Compressed region %zu starts past entry end (%" PRIu64 " pages)\n", + index, pe->nr_pages); + return -1; + } + pages_left = pe->nr_pages - first_page; + *block_pages = pages_left < region_pages ? + (unsigned int)pages_left : region_pages; + *block_bytes = (size_t)*block_pages * PAGE_SIZE; + return 0; +} + +static int compressed_block_storage(const PagemapEntry *pe, size_t index, + enum restore_vma_io_storage *storage, + unsigned int *block_pages) +{ + size_t block_bytes; + uint32_t compressed_size; + + if (compressed_block_layout(pe, index, block_pages, &block_bytes)) + return -1; + + compressed_size = pe->compressed_size[index]; + if (!compressed_size) + *storage = VMA_IO_ZERO; + else if (compressed_size == block_bytes) + *storage = VMA_IO_PACKED_RAW; + else + *storage = VMA_IO_ENCODED; + return 0; +} + static int piov_add_compressed_blocks(struct page_read_iov *piov, struct page_read *pr, - unsigned long nr_pages) + unsigned long nr_pages, + enum restore_vma_io_storage storage) { - unsigned int region_pages = 0; + unsigned int region_pages = pagemap_region_pages(pr->pe); size_t first = piov->n_compressed_size; unsigned long n_blocks; + uint16_t *new_bp = NULL; uint32_t *new_cs; uint64_t added = 0; size_t new_n; unsigned long i; - if (pr->pe->has_region_pages && pr->pe->region_pages) - region_pages = pr->pe->region_pages; - if (region_pages) { uint64_t pages_consumed = (uint64_t)pr->compressed_size_index * region_pages; uint64_t end_page = pages_consumed + nr_pages; @@ -493,30 +630,33 @@ static int piov_add_compressed_blocks(struct page_read_iov *piov, return -1; piov->compressed_size = new_cs; - for (i = 0; i < n_blocks; i++) { - uint32_t cs = pr->pe->compressed_size[pr->compressed_size_index + i]; - - piov->compressed_size[first + i] = cs; - added += cs; - } - if (region_pages) { - uint64_t pages_consumed_at_start = (uint64_t)pr->compressed_size_index * region_pages; - uint16_t *new_bp = xrealloc(piov->block_pages, new_n * sizeof(*piov->block_pages)); - + new_bp = xrealloc(piov->block_pages, new_n * sizeof(*piov->block_pages)); if (!new_bp) return -1; piov->block_pages = new_bp; + } - for (i = 0; i < n_blocks; i++) { - uint64_t pages_so_far = pages_consumed_at_start + i * region_pages; - uint64_t pages_left = pr->pe->nr_pages - pages_so_far; - uint16_t block_pages = (uint16_t)region_pages; + for (i = 0; i < n_blocks; i++) { + size_t idx = pr->compressed_size_index + i; + enum restore_vma_io_storage block_storage; + unsigned int block_pages; + uint32_t cs = pr->pe->compressed_size[idx]; - if (pages_left < region_pages) - block_pages = (uint16_t)pages_left; - piov->block_pages[first + i] = block_pages; + if (compressed_block_storage(pr->pe, idx, &block_storage, + &block_pages)) + return -1; + if ((storage == VMA_IO_PACKED_RAW || storage == VMA_IO_ZERO) && + block_storage != storage) { + pr_err("Block %zu has storage kind %d, expected %d\n", + idx, block_storage, storage); + return -1; } + + piov->compressed_size[first + i] = cs; + if (region_pages) + piov->block_pages[first + i] = (uint16_t)block_pages; + added += cs; } piov->total_compressed_size += added; @@ -527,11 +667,28 @@ static int piov_add_compressed_blocks(struct page_read_iov *piov, return 0; } -static int enqueue_async_iov(struct page_read *pr, void *buf, unsigned long len, struct list_head *to) +static int enqueue_async_iov(struct page_read *pr, void *buf, unsigned long len, + struct list_head *to, + enum restore_vma_io_storage storage) { struct page_read_iov *pr_iov; struct iovec *iov; + if (!len || len % PAGE_SIZE) { + pr_err("Invalid async page-read length %lu\n", len); + return -1; + } + if (storage == VMA_IO_UNCOMPRESSED && pr->pe && + pr->pe->n_compressed_size) { + pr_err("Uncompressed async job has compression metadata\n"); + return -1; + } + if (storage != VMA_IO_UNCOMPRESSED && + (!pr->pe || !pr->pe->n_compressed_size)) { + pr_err("Packed async job has no compression metadata\n"); + return -1; + } + pr_iov = xzalloc(sizeof(*pr_iov)); if (!pr_iov) return -1; @@ -551,6 +708,7 @@ static int enqueue_async_iov(struct page_read *pr, void *buf, unsigned long len, pr_iov->to = iov; pr_iov->nr = 1; pr_iov->n_pages = len / PAGE_SIZE; + pr_iov->storage = storage; /* * For uncompressed entries, the end offset is simply @@ -559,19 +717,21 @@ static int enqueue_async_iov(struct page_read *pr, void *buf, unsigned long len, * region mode) so that process_async_reads() knows how much * compressed data to read from the image. */ - if (!pr->pe || !pr->pe->n_compressed_size) { + if (storage == VMA_IO_UNCOMPRESSED) pr_iov->end += len; - } else if (piov_add_compressed_blocks(pr_iov, pr, len / PAGE_SIZE)) { - xfree(pr_iov->compressed_size); - xfree(pr_iov->block_pages); - xfree(iov); - xfree(pr_iov); - return -1; - } + else if (piov_add_compressed_blocks(pr_iov, pr, len / PAGE_SIZE, storage)) + goto err; list_add_tail(&pr_iov->l, to); return 0; + +err: + xfree(pr_iov->block_pages); + xfree(pr_iov->compressed_size); + xfree(iov); + xfree(pr_iov); + return -1; } int pagemap_render_iovec(struct list_head *from, struct task_restore_args *ta) @@ -632,6 +792,7 @@ int pagemap_render_iovec(struct list_head *from, struct task_restore_args *ta) rio->nr_iovs = piov->nr; rio->off = piov->from; + rio->storage = piov->storage; /* Ordinary reads do not carry compressed block metadata. */ if (!piov->n_compressed_size) { @@ -723,26 +884,94 @@ static int advance_compressed_offsets(struct page_read *pr, unsigned long nr) return 0; } -static unsigned long compressed_async_chunk_pages(struct page_read *pr, unsigned long pages_left) +static int compressed_request_has_lz4(struct page_read *pr, + unsigned long nr_pages, bool *has_lz4, + bool *mixed_direct) { - unsigned int region_pages = 0; - unsigned long chunk = pages_left; + size_t index = pr->compressed_size_index; + enum restore_vma_io_storage previous_storage = VMA_IO_UNCOMPRESSED; + bool have_previous_storage = false; + unsigned long pages_done = 0; - if (pr->pe && pr->pe->has_region_pages && pr->pe->region_pages) - region_pages = pr->pe->region_pages; - if (chunk > ASYNC_BATCH_MAX_PAGES) - chunk = ASYNC_BATCH_MAX_PAGES; + *has_lz4 = false; + *mixed_direct = false; + while (pages_done < nr_pages) { + enum restore_vma_io_storage storage; + unsigned int block_pages; - if (region_pages && chunk < pages_left) { - chunk -= chunk % region_pages; - if (!chunk) - chunk = region_pages; + if (compressed_block_storage(pr->pe, index, &storage, + &block_pages)) + return -1; + if (block_pages > nr_pages - pages_done) { + pr_err("Compressed request ends inside block %zu\n", index); + return -1; + } + if (storage == VMA_IO_ENCODED) + *has_lz4 = true; + else if (have_previous_storage && storage != previous_storage) + *mixed_direct = true; + if (storage != VMA_IO_ENCODED) { + previous_storage = storage; + have_previous_storage = true; + } + pages_done += block_pages; + index++; } - return chunk; + return 0; } -static int pagemap_enqueue_iovec_one(struct page_read *pr, void *buf, unsigned long len, struct list_head *to) +static int compressed_storage_run(struct page_read *pr, + unsigned long pages_left, + enum restore_vma_io_storage *storage, + unsigned long *run_pages, + size_t *run_blocks) +{ + size_t index = pr->compressed_size_index; + enum restore_vma_io_storage first_storage; + unsigned int block_pages; + + if (compressed_block_storage(pr->pe, index, &first_storage, + &block_pages)) + return -1; + + *run_pages = 0; + *run_blocks = 0; + while (*run_pages < pages_left) { + enum restore_vma_io_storage block_storage = first_storage; + + if (*run_pages) { + if (compressed_block_storage(pr->pe, index, &block_storage, &block_pages)) + return -1; + } + if (block_storage != first_storage) + break; + if (block_pages > pages_left - *run_pages) { + pr_err("Compressed request ends inside block %zu\n", index); + return -1; + } + if (*run_pages && + block_pages > ASYNC_BATCH_MAX_PAGES - *run_pages) + break; + *run_pages += block_pages; + (*run_blocks)++; + index++; + if (*run_pages >= ASYNC_BATCH_MAX_PAGES) + break; + } + + if (!*run_pages) { + pr_err("Unable to form compressed storage run at block %zu\n", + pr->compressed_size_index); + return -1; + } + *storage = first_storage; + return 0; +} + +static int pagemap_enqueue_iovec_one(struct page_read *pr, void *buf, + unsigned long len, struct list_head *to, + enum restore_vma_io_storage storage) { struct page_read_iov *cur_async = NULL; struct iovec *iov; @@ -750,8 +979,8 @@ static int pagemap_enqueue_iovec_one(struct page_read *pr, void *buf, unsigned l bool added_iov = false; unsigned int new_region_pages = 0; - if (pr->pe && pr->pe->has_region_pages && pr->pe->region_pages) - new_region_pages = pr->pe->region_pages; + if (storage != VMA_IO_UNCOMPRESSED) + new_region_pages = pagemap_region_pages(pr->pe); if (!list_empty(to)) cur_async = list_entry(to->prev, struct page_read_iov, l); @@ -763,30 +992,33 @@ static int pagemap_enqueue_iovec_one(struct page_read *pr, void *buf, unsigned l * Start the new preadv request here. */ if (!cur_async || pr->pi_off != cur_async->end) - return enqueue_async_iov(pr, buf, len, to); + return enqueue_async_iov(pr, buf, len, to, storage); + + /* Different storage kinds have different restore and dedup rules. */ + if (cur_async->storage != storage) + return enqueue_async_iov(pr, buf, len, to, storage); /* * Don't merge piovs with different region modes: each piov is * decompressed with a single algorithm. */ if (cur_async->region_pages != new_region_pages) - return enqueue_async_iov(pr, buf, len, to); + return enqueue_async_iov(pr, buf, len, to, storage); /* * Cap a compressed async batch by its UNCOMPRESSED page count. - * process_async_reads() stages the whole batch's compressed payload in - * one buffer, while the decompressed destination spans - * n_pages * PAGE_SIZE. Bounding pages bounds both because compressed - * payload never exceeds its decoded size. A compressed-byte cap alone - * would not bound the destination for highly compressible data (or - * all-zero pages, whose compressed size is 0 and never trips a byte - * cap), so a batch could grow to many GiB and exhaust host RAM. + * process_async_reads() reads an encoded batch into one buffer and + * allocates per-block job metadata, while delayed raw/zero batches copy + * their block metadata into restorer memory. Bounding pages bounds both + * working sets. A compressed-byte cap alone would not bound metadata for + * highly compressible or zero-filled data, whose payload can be tiny or + * empty while the logical range grows to many GiB. * The check is gated on compression; uncompressed async reads go * straight into their destination iovecs and need no cap. */ - if (pr->pe && pr->pe->n_compressed_size && + if (storage != VMA_IO_UNCOMPRESSED && cur_async->n_pages + len / PAGE_SIZE > ASYNC_BATCH_MAX_PAGES) - return enqueue_async_iov(pr, buf, len, to); + return enqueue_async_iov(pr, buf, len, to, storage); /* * This read is pure continuation of the previous one. Let's @@ -802,7 +1034,7 @@ static int pagemap_enqueue_iovec_one(struct page_read *pr, void *buf, unsigned l unsigned int n_iovs = cur_async->nr + 1; if (n_iovs > IOV_MAX) - return enqueue_async_iov(pr, buf, len, to); + return enqueue_async_iov(pr, buf, len, to, storage); iov = xrealloc(cur_async->to, n_iovs * sizeof(*iov)); if (!iov) @@ -828,11 +1060,10 @@ static int pagemap_enqueue_iovec_one(struct page_read *pr, void *buf, unsigned l /* Extend the end offset. For compressed entries, append * per-block sizes and advance by compressed bytes. */ - if (!pr->pe || !pr->pe->n_compressed_size) { + if (storage == VMA_IO_UNCOMPRESSED) { cur_async->end += len; - } else if (piov_add_compressed_blocks(cur_async, pr, len / PAGE_SIZE)) { + } else if (piov_add_compressed_blocks(cur_async, pr, len / PAGE_SIZE, storage)) goto rollback_iov; - } return 0; @@ -855,15 +1086,29 @@ int pagemap_enqueue_iovec(struct page_read *pr, void *buf, unsigned long len, st unsigned int original_nr = 0; size_t original_last_iov_len = 0; unsigned long nr_pages = len / PAGE_SIZE; + bool has_lz4; + bool mixed_direct; + bool parallel_zero = false; off_t pi_off; size_t compressed_size_index; unsigned int region_block_offset; unsigned long done = 0; int ret = 0; - if (!pr->pe || !pr->pe->n_compressed_size || - nr_pages <= ASYNC_BATCH_MAX_PAGES) - return pagemap_enqueue_iovec_one(pr, buf, len, to); + if (!len || len % PAGE_SIZE) { + pr_err("Invalid pagemap I/O length %lu\n", len); + return -1; + } + + if (!pr->pe || !pr->pe->n_compressed_size) + return pagemap_enqueue_iovec_one(pr, buf, len, to, + VMA_IO_UNCOMPRESSED); + + if (compressed_request_has_lz4(pr, nr_pages, &has_lz4, + &mixed_direct)) + return -1; + if (to == &pr->async) + parallel_zero = compressed_restore_has_parallel_capacity(opts.decompress_threads); pi_off = pr->pi_off; compressed_size_index = pr->compressed_size_index; @@ -872,17 +1117,43 @@ int pagemap_enqueue_iovec(struct page_read *pr, void *buf, unsigned long len, st original_tail = list_entry(to->prev, struct page_read_iov, l); original_end = original_tail->end; original_n_compressed_size = original_tail->n_compressed_size; - original_total_compressed_size = original_tail->total_compressed_size; + original_total_compressed_size = + original_tail->total_compressed_size; original_n_pages = original_tail->n_pages; original_nr = original_tail->nr; original_last_iov_len = original_tail->to[original_nr - 1].iov_len; } while (done < nr_pages) { - unsigned long chunk = compressed_async_chunk_pages(pr, nr_pages - done); + enum restore_vma_io_storage storage; + unsigned long chunk; + size_t blocks; + + if (compressed_storage_run(pr, nr_pages - done, &storage, + &chunk, &blocks)) { + ret = -1; + break; + } + + /* Premapped large zero runs can share the decompression worker pool. */ + if (storage == VMA_IO_ZERO && parallel_zero && blocks > 1 && + chunk * PAGE_SIZE >= PARALLEL_RESTORE_MIN_BATCH_BYTES) + storage = VMA_IO_ENCODED; + + /* + * Keep short raw/zero islands in an in-process encoded batch when + * the request either contains LZ4 or was premapped because it has + * too many direct runs. Adjacent short runs then coalesce, while + * long raw/zero runs retain their direct paths. PIE requests remain + * split because PIE cannot decode an encoded batch. + */ + if (storage != VMA_IO_ENCODED && + (has_lz4 || (mixed_direct && to == &pr->async)) && + chunk * PAGE_SIZE < DIRECT_COMPRESSED_RUN_MIN_BYTES) + storage = VMA_IO_ENCODED; ret = pagemap_enqueue_iovec_one(pr, (char *)buf + done * PAGE_SIZE, - chunk * PAGE_SIZE, to); + chunk * PAGE_SIZE, to, storage); if (ret) break; @@ -898,13 +1169,16 @@ int pagemap_enqueue_iovec(struct page_read *pr, void *buf, unsigned long len, st pr->region_block_offset = region_block_offset; if (ret) { /* - * One logical request may span several batches. If a later - * allocation or metadata check fails, remove each new piov and - * restore any pre-existing tail extended by an earlier chunk. + * One logical request may be split into multiple storage runs. If + * a later allocation or metadata check fails, remove every new + * piov and restore a pre-existing tail that earlier runs extended. + * Leaving a prefix queued would make the caller's error cleanup hit + * close_page_read() with an unexpected non-empty async list. */ while (!list_empty(to) && (!original_tail || to->prev != &original_tail->l)) { - struct page_read_iov *piov = list_entry(to->prev, struct page_read_iov, l); + struct page_read_iov *piov = + list_entry(to->prev, struct page_read_iov, l); list_del(&piov->l); xfree(piov->compressed_size); @@ -914,17 +1188,60 @@ int pagemap_enqueue_iovec(struct page_read *pr, void *buf, unsigned long len, st } if (original_tail) { original_tail->end = original_end; - original_tail->n_compressed_size = original_n_compressed_size; - original_tail->total_compressed_size = original_total_compressed_size; + original_tail->n_compressed_size = + original_n_compressed_size; + original_tail->total_compressed_size = + original_total_compressed_size; original_tail->n_pages = original_n_pages; original_tail->nr = original_nr; - original_tail->to[original_nr - 1].iov_len = original_last_iov_len; + original_tail->to[original_nr - 1].iov_len = + original_last_iov_len; } } return ret; } +bool pagemap_iovec_is_direct_compatible(const struct list_head *from) +{ + const struct page_read_iov *piov; + bool has_file_backed = false; + + list_for_each_entry(piov, from, l) { + uint64_t destination_bytes = 0; + unsigned int i; + + if (!piov->nr || piov->storage == VMA_IO_ENCODED) + return false; + if (piov->storage != VMA_IO_UNCOMPRESSED && + piov->storage != VMA_IO_PACKED_RAW && + piov->storage != VMA_IO_ZERO) + return false; + if (piov->storage != VMA_IO_ZERO && + (piov->from < 0 || piov->end < piov->from || + piov->from % (off_t)PAGE_SIZE || + (piov->end - piov->from) % (off_t)PAGE_SIZE)) + return false; + if (piov->storage != VMA_IO_ZERO) + has_file_backed = true; + + for (i = 0; i < piov->nr; i++) { + if ((uintptr_t)piov->to[i].iov_base % PAGE_SIZE || + piov->to[i].iov_len % PAGE_SIZE) + return false; + if (destination_bytes > + UINT64_MAX - piov->to[i].iov_len) + return false; + destination_bytes += piov->to[i].iov_len; + } + if (piov->storage != VMA_IO_ZERO && + destination_bytes != (uint64_t)(piov->end - piov->from)) + return false; + } + + return has_file_backed; +} + static int maybe_read_page_local(struct page_read *pr, unsigned long vaddr, unsigned long nr, void *buf, unsigned flags) { int ret; @@ -972,6 +1289,121 @@ static int pread_full(int fd, void *buf, size_t count, off_t offset) return 0; } +/* Read without logging: an error is reported only if this payload is used. */ +static void encoded_prefetch_read(void *arg) +{ + struct encoded_prefetch *prefetch = arg; + + while (prefetch->done < prefetch->count) { + ssize_t ret; + + ret = pread(prefetch->fd, prefetch->buffer + prefetch->done, + prefetch->count - prefetch->done, + prefetch->offset + (off_t)prefetch->done); + if (ret < 0) { + if (errno == EINTR) + continue; + prefetch->saved_errno = errno; + break; + } + if (!ret) + break; + prefetch->done += ret; + } + prefetch->complete = true; +} + +static void encoded_prefetch_disable(struct encoded_read_ctx *ctx) +{ + ctx->prefetched_piov = NULL; + memset(&ctx->prefetch, 0, sizeof(ctx->prefetch)); + xfree(ctx->prefetch_buffer); + ctx->prefetch_buffer = NULL; + ctx->prefetch_cap = 0; + if (ctx->prefetch_batch_acquired) { + ctx->prefetch_batch_acquired = false; + decompression_batch_release(); + } +} + +static bool encoded_prefetch_prepare(struct encoded_read_ctx *ctx, int fd, + off_t offset, size_t count) +{ + void *new_buffer; + + BUG_ON(ctx->prefetched_piov || ctx->prefetch.complete || !count); + if (!ctx->prefetch_batch_acquired) { + if (!decompression_batch_try_acquire()) + return false; + ctx->prefetch_batch_acquired = true; + } + if (ctx->prefetch_cap < count) { + new_buffer = xrealloc(ctx->prefetch_buffer, count); + if (!new_buffer) { + encoded_prefetch_disable(ctx); + return false; + } + ctx->prefetch_buffer = new_buffer; + ctx->prefetch_cap = count; + } + + ctx->prefetch.buffer = ctx->prefetch_buffer; + ctx->prefetch.count = count; + ctx->prefetch.done = 0; + ctx->prefetch.offset = offset; + ctx->prefetch.fd = fd; + ctx->prefetch.saved_errno = 0; + return true; +} + +static void encoded_prefetch_publish(struct encoded_read_ctx *ctx, + struct page_read_iov *piov) +{ + void *buffer; + size_t capacity; + + BUG_ON(!ctx->prefetch.complete); + buffer = ctx->compressed; + capacity = ctx->compressed_cap; + ctx->compressed = ctx->prefetch_buffer; + ctx->compressed_cap = ctx->prefetch_cap; + ctx->prefetch_buffer = buffer; + ctx->prefetch_cap = capacity; + ctx->prefetched_piov = piov; +} + +static int encoded_prefetch_take(struct encoded_read_ctx *ctx, + struct page_read_iov *piov) +{ + if (!ctx->prefetched_piov) + return 0; + if (ctx->prefetched_piov != piov) { + pr_err("Encoded prefetch does not match the next request\n"); + return -1; + } + ctx->prefetched_piov = NULL; + ctx->prefetch.complete = false; + if (ctx->prefetch.count != piov->total_compressed_size) { + pr_err("Encoded prefetch size changed before use\n"); + return -1; + } + + if (ctx->prefetch.saved_errno) { + errno = ctx->prefetch.saved_errno; + pr_perror("Can't read %zu encoded bytes at offset %lld", + ctx->prefetch.count, (long long)ctx->prefetch.offset); + return -1; + } + if (ctx->prefetch.done != ctx->prefetch.count) { + pr_err("Short encoded read %zu/%zu at offset %lld\n", + ctx->prefetch.done, ctx->prefetch.count, + (long long)ctx->prefetch.offset); + return -1; + } + + return 1; +} + /* * Advance compressed offset tracking by @nr pages without * reading any data. Used when enqueuing async compressed reads. @@ -1091,15 +1523,10 @@ static int read_compressed_pages_region(struct page_read *pr, int fd, unsigned long vaddr, unsigned long nr, void *buf) { - unsigned int region_pages = pr->pe->region_pages; - size_t region_bytes_max = (size_t)region_pages * PAGE_SIZE; unsigned long pages_done = 0; + size_t scratch_cap = 0; int rc = -1; - char *scratch; - - scratch = xmalloc(region_bytes_max); - if (!scratch) - return -1; + char *scratch = NULL; while (pages_done < nr) { size_t idx = pr->compressed_size_index; @@ -1146,7 +1573,16 @@ static int read_compressed_pages_region(struct page_read *pr, int fd, pr->pi_off + (off_t)off * PAGE_SIZE)) goto out; } else if (off == 0 && take == this_region) { + char *new_scratch; + /* Whole region in one go: decompress straight into the dest. */ + if (scratch_cap < cs) { + new_scratch = xrealloc(scratch, cs); + if (!new_scratch) + goto out; + scratch = new_scratch; + scratch_cap = cs; + } if (pread_full(fd, scratch, cs, pr->pi_off)) goto out; if (decompress_region(scratch, cs, this_region, @@ -1156,15 +1592,25 @@ static int read_compressed_pages_region(struct page_read *pr, int fd, goto out; } } else { + char *new_scratch; + /* * Partial slice: keep the decompressed region around * because incremental restores can read alternating * pages from the same parent region. */ - if (!region_cache_hit(pr, region_vaddr, this_bytes) && - region_cache_load(pr, fd, scratch, region_vaddr, idx, - cs, this_region, this_bytes)) - goto out; + if (!region_cache_hit(pr, region_vaddr, this_bytes)) { + if (scratch_cap < cs) { + new_scratch = xrealloc(scratch, cs); + if (!new_scratch) + goto out; + scratch = new_scratch; + scratch_cap = cs; + } + if (region_cache_load(pr, fd, scratch, region_vaddr, idx, + cs, this_region, this_bytes)) + goto out; + } memcpy((char *)buf + pages_done * PAGE_SIZE, pr->cached_region + (size_t)off * PAGE_SIZE, out_bytes); @@ -1352,22 +1798,30 @@ static int maybe_read_page_img_streamer(struct page_read *pr, unsigned long vadd } /* - * Streaming + compressed reader: reads compressed data sequentially - * from the pipe (no seeking), decompresses page-by-page. + * Restore one compressed pagemap extent from a forward-only image stream. + * + * Unlike the local-image path, this reader cannot pread individual blocks or + * retry them out of order. It therefore validates each bounded metadata batch + * before consuming the corresponding packed payload, reads that payload once, + * and keeps it alive until every LZ4 worker has completed. Zero and raw pages + * are materialized inline; only LZ4-compressed pages enter the persistent worker + * pool. Reader offsets and the compressed-size cursor are committed only after + * the entire requested extent is decoded; io_complete runs before publication. + * + * Entries whose metadata was elided because every page was stored raw retain + * the ordinary streaming representation and use maybe_read_page_img_streamer(). */ static int maybe_read_page_img_streamer_compressed(struct page_read *pr, unsigned long vaddr, unsigned long nr, void *buf, unsigned flags) { + struct page_read *owner = pr->encoded_read_owner; + struct encoded_read_ctx *ctx = NULL; + uint64_t total_payload = 0; + unsigned long pages_done = 0; + ssize_t bytes; int fd; - ssize_t ret = 0; - size_t curr = 0; - off_t compressed_offset = 0; - char compressed_buf[PAGE_COMPRESSED_SIZE_BOUND]; - unsigned long i; + int ret = -1; - /* - * Fall back to uncompressed streamer for entries without compressed data. - * Region mode is not supported in streaming restore (checked in config). - */ + /* Select the representation before consuming alignment or payload bytes. */ if (!pr->pe->n_compressed_size) return maybe_read_page_img_streamer(pr, vaddr, nr, buf, flags); @@ -1385,18 +1839,18 @@ static int maybe_read_page_img_streamer_compressed(struct page_read *pr, unsigne char discard[PAGE_PADDING_CHUNK]; size_t chunk = min(pr->stream_padding, sizeof(discard)); - ret = read(fd, discard, chunk); - if (ret == 0) { + bytes = read(fd, discard, chunk); + if (bytes == 0) { pr_err("Reached EOF while skipping compressed page-image alignment\n"); return -1; } - if (ret < 0) { + if (bytes < 0) { if (errno == EINTR) continue; pr_perror("Can't skip compressed page-image alignment"); return -1; } - pr->stream_padding -= ret; + pr->stream_padding -= bytes; } pr_debug("\tpr%lu-%u Read page from self %lx/%" PRIx64 "\n", @@ -1404,90 +1858,159 @@ static int maybe_read_page_img_streamer_compressed(struct page_read *pr, unsigne BUG_ON(pr->cvaddr != vaddr); - if (pr->compressed_size_index + nr > pr->pe->n_compressed_size) { + if (pr->compressed_size_index > pr->pe->n_compressed_size || + nr > pr->pe->n_compressed_size - pr->compressed_size_index) { pr_err("Compressed size index out of bounds: %zu + %lu > %zu\n", pr->compressed_size_index, nr, pr->pe->n_compressed_size); return -1; } - - for (i = 0; i < nr; i++) { - size_t idx = pr->compressed_size_index + i; - uint32_t cs = pr->pe->compressed_size[idx]; - size_t rd = 0; - - if (cs > PAGE_COMPRESSED_SIZE_BOUND) { - pr_err("Invalid compressed size %u for page %zu\n", - cs, idx); - return -1; - } - - if (cs == 0) { - /* Zero page, nothing stored in the image */ - memset(buf + curr, 0, PAGE_SIZE); - curr += PAGE_SIZE; - continue; - } - - if (cs == PAGE_SIZE) { - /* Stored raw, read directly into output */ - while (rd < PAGE_SIZE) { - ret = read(fd, buf + curr + rd, PAGE_SIZE - rd); - if (ret == 0) { - pr_err("Reached EOF reading raw page\n"); - return -1; - } else if (ret < 0) { - if (errno == EINTR) - continue; - pr_perror("Can't read raw page"); - return -1; - } - rd += ret; - } - } else { - while (rd < cs) { - /* - * cs is bounded by PAGE_COMPRESSED_SIZE_BOUND - * (validated above), so cs - rd fits in the - * stack buffer. Re-state the bound explicitly - * so the compiler's fortified-read analysis - * does not warn about a possible overflow into - * compressed_buf. - */ - size_t to_read = cs - rd; - - if (to_read > sizeof(compressed_buf) - rd) - to_read = sizeof(compressed_buf) - rd; - - ret = read(fd, compressed_buf + rd, to_read); - if (ret == 0) { - pr_err("Reached EOF reading compressed page\n"); - return -1; - } else if (ret < 0) { - if (errno == EINTR) - continue; - pr_perror("Can't read compressed page"); - return -1; - } - rd += ret; - } - - if (decompress_data(compressed_buf, cs, PAGE_SIZE, buf + curr)) { - pr_err("Decompression failed for page %zu\n", idx); - return -1; - } - } - - curr += PAGE_SIZE; - compressed_offset += cs; + if (pr->pi_off < 0) { + pr_err("Invalid negative compressed stream offset: %jd\n", + (intmax_t)pr->pi_off); + return -1; } + while (pages_done < nr) { + unsigned long batch_pages = min(nr - pages_done, + ASYNC_BATCH_MAX_PAGES); + size_t first = pr->compressed_size_index + pages_done; + size_t batch_payload = 0; + size_t payload_offset = 0; + size_t jobs_uncompressed = 0; + size_t nr_jobs = 0; + size_t i; + + /* Validate and size the batch before consuming any bytes. */ + for (i = 0; i < batch_pages; i++) { + uint32_t cs = pr->pe->compressed_size[first + i]; + + if (cs > PAGE_SIZE) { + pr_err("Invalid compressed size %u for page %zu\n", + cs, first + i); + goto out; + } + if (batch_payload > ASYNC_BATCH_MAX_BYTES - cs) { + pr_err("Compressed streaming batch exceeds %lu bytes\n", + ASYNC_BATCH_MAX_BYTES); + goto out; + } + batch_payload += cs; + if (cs && cs < PAGE_SIZE) + nr_jobs++; + } + + if (total_payload > (uint64_t)OFF_MAX - (uint64_t)pr->pi_off || + batch_payload > (uint64_t)OFF_MAX - (uint64_t)pr->pi_off - + total_payload) { + pr_err("Compressed stream offset overflows at page %zu\n", first); + goto out; + } + + /* Lazily reserve one working set and reuse it for this active read. */ + if (batch_payload || nr_jobs) { + if (!owner) { + pr_err("Streaming page reader has no encoded-read context owner\n"); + goto out; + } + if (!owner->encoded_read_ctx) + owner->encoded_read_ctx = + xzalloc(sizeof(*owner->encoded_read_ctx)); + if (!owner->encoded_read_ctx) + goto out; + ctx = owner->encoded_read_ctx; + + if (!ctx->batch_acquired) + encoded_read_ctx_begin_work(ctx); + if (ctx->compressed_cap < batch_payload) { + void *new_compressed = + xrealloc(ctx->compressed, batch_payload); + + if (!new_compressed) + goto out; + ctx->compressed = new_compressed; + ctx->compressed_cap = batch_payload; + } + if (ctx->jobs_cap < nr_jobs) { + void *new_jobs = xrealloc(ctx->jobs, + nr_jobs * sizeof(*ctx->jobs)); + + if (!new_jobs) + goto out; + ctx->jobs = new_jobs; + ctx->jobs_cap = nr_jobs; + } + } + + /* Consume this batch in one sequential read from the image stream. */ + if (batch_payload) { + bytes = read_all(fd, ctx->compressed, batch_payload); + if (bytes < 0) { + pr_perror("Can't read compressed streaming batch"); + goto out; + } + if ((size_t)bytes != batch_payload) { + pr_err("Reached EOF reading compressed streaming batch: %zd of %zu bytes\n", + bytes, batch_payload); + goto out; + } + } + + /* Fill zero/raw pages and describe disjoint LZ4 destinations. */ + nr_jobs = 0; + for (i = 0; i < batch_pages; i++) { + size_t idx = first + i; + uint32_t cs = pr->pe->compressed_size[idx]; + char *dst = (char *)buf + + (pages_done + i) * PAGE_SIZE; + + if (!cs) { + memset(dst, 0, PAGE_SIZE); + } else if (cs == PAGE_SIZE) { + memcpy(dst, ctx->compressed + payload_offset, + PAGE_SIZE); + } else { + struct decompress_job *job = &ctx->jobs[nr_jobs++]; + + job->src = ctx->compressed + payload_offset; + job->dst = dst; + job->compressed_size = cs; + job->pages = 1; + job->block_index = i; + jobs_uncompressed += PAGE_SIZE; + } + payload_offset += cs; + } + + if (payload_offset != batch_payload) { + pr_err("Compressed streaming batch metadata mismatch: %zu != %zu\n", + payload_offset, batch_payload); + goto out; + } + /* Source pointers remain stable until all jobs in the batch finish. */ + if (nr_jobs && + decompress_jobs_parallel_pool(&ctx->pool, ctx->jobs, nr_jobs, + jobs_uncompressed, + opts.decompress_threads)) { + pr_err("Unable to decompress streaming batch at page %zu\n", + first); + goto out; + } + + total_payload += batch_payload; + pages_done += batch_pages; + } + + /* Publish completion before advancing externally visible stream cursors. */ + ret = 0; if (pr->io_complete) ret = pr->io_complete(pr, vaddr, nr); - pr->pi_off += compressed_offset; + pr->pi_off += (off_t)total_payload; pr->compressed_size_index += nr; +out: + encoded_read_ctx_end_work(ctx); return ret; } @@ -1553,24 +2076,31 @@ static void advance_piov(struct page_read_iov *piov, ssize_t len) { ssize_t olen = len; int onr = piov->nr; + unsigned int consumed = 0; + piov->from += len; - while (len) { - struct iovec *cur = piov->to; - - if (cur->iov_len <= len) { - piov->to++; - piov->nr--; - len -= cur->iov_len; - continue; - } - - cur->iov_base += len; - cur->iov_len -= len; - break; + while (consumed < piov->nr && + piov->to[consumed].iov_len <= (size_t)len) { + len -= piov->to[consumed].iov_len; + consumed++; } - pr_debug("Advanced iov %zu bytes, %d->%d iovs, %zu tail\n", olen, onr, piov->nr, len); + if (consumed) { + piov->nr -= consumed; + if (piov->nr) + memmove(piov->to, piov->to + consumed, + piov->nr * sizeof(*piov->to)); + } + if (len) { + BUG_ON(!piov->nr); + piov->to[0].iov_base = + (char *)piov->to[0].iov_base + len; + piov->to[0].iov_len -= len; + } + + pr_debug("Advanced iov %zd bytes, %d->%d iovs, %zd tail\n", + olen, onr, piov->nr, len); } /* @@ -1594,205 +2124,404 @@ static void drain_async_queue(struct page_read *pr) } /* - * Compressed async read: read all compressed - * data in one pread(), then decompress - * block-by-block into the destination iovecs. - * In per-page mode each block is one page; in - * region mode each block is up to region_pages - * pages. + * Validate a raw/zero compressed-format request before it bypasses LZ4. + * Block metadata must describe exactly the destination bytes; packed-raw + * payload bytes must also match the pages-image extent, while zero requests + * have no payload. This keeps both the local preadv path and rendered PIE + * requests from trusting inconsistent counts. */ -static int process_compressed_async_read(struct page_read *pr, int fd, - struct page_read_iov *piov) +static int validate_direct_compressed_iov(const struct page_read_iov *piov) { - char *comp_buf; - off_t comp_off; - size_t total = piov->total_compressed_size; - int iov_idx = 0; - size_t iov_off = 0; - char *region_scratch = NULL; - size_t region_scratch_cap = 0; + uint64_t payload_bytes = 0; + uint64_t output_bytes = 0; + uint64_t iov_bytes = 0; + size_t i; + + if (piov->storage != VMA_IO_PACKED_RAW && + piov->storage != VMA_IO_ZERO) + return -1; + if (!piov->n_compressed_size || !piov->compressed_size) { + pr_err("Direct compressed I/O job has no block metadata\n"); + return -1; + } + if (piov->region_pages && !piov->block_pages) { + pr_err("Direct region I/O job has no block-page metadata\n"); + return -1; + } + + /* Sum and validate each block without overflowing either byte count. */ + for (i = 0; i < piov->n_compressed_size; i++) { + unsigned int block_pages = piov->region_pages ? + piov->block_pages[i] : 1; + uint64_t block_bytes; + uint32_t compressed_size = piov->compressed_size[i]; + + if (!block_pages || + (piov->region_pages && block_pages > piov->region_pages)) { + pr_err("Invalid page count %u for direct block %zu\n", + block_pages, i); + return -1; + } + block_bytes = (uint64_t)block_pages * PAGE_SIZE; + if (piov->storage == VMA_IO_PACKED_RAW && + compressed_size != block_bytes) { + pr_err("Packed-raw block %zu has size %u, expected %" PRIu64 "\n", + i, compressed_size, block_bytes); + return -1; + } + if (piov->storage == VMA_IO_ZERO && compressed_size) { + pr_err("Zero block %zu has payload size %u\n", i, + compressed_size); + return -1; + } + if (payload_bytes > UINT64_MAX - compressed_size || + output_bytes > UINT64_MAX - block_bytes) { + pr_err("Direct compressed I/O size overflows\n"); + return -1; + } + payload_bytes += compressed_size; + output_bytes += block_bytes; + } + + /* Independently account the scatter/gather destination. */ + for (i = 0; i < piov->nr; i++) { + if (iov_bytes > UINT64_MAX - piov->to[i].iov_len) { + pr_err("Direct compressed destination size overflows\n"); + return -1; + } + iov_bytes += piov->to[i].iov_len; + } + + if (payload_bytes != piov->total_compressed_size || + output_bytes != iov_bytes || piov->end < piov->from || + payload_bytes != (uint64_t)(piov->end - piov->from) || + piov->n_pages > UINT64_MAX / PAGE_SIZE || + output_bytes != (uint64_t)piov->n_pages * PAGE_SIZE) { + pr_err("Inconsistent direct compressed I/O sizes: payload=%" PRIu64 + " metadata=%" PRIu64 " output=%" PRIu64 " iov=%" PRIu64 + " range=%jd\n", payload_bytes, + piov->total_compressed_size, output_bytes, iov_bytes, + (intmax_t)(piov->end - piov->from)); + return -1; + } + + return 0; +} + +static int transfer_async_block(const struct page_read_iov *piov, + unsigned int *iov_index, size_t *iov_offset, + const char *src, size_t bytes) +{ + while (bytes) { + char *dst; + size_t chunk; + + while (*iov_index < piov->nr && + *iov_offset >= piov->to[*iov_index].iov_len) { + *iov_offset -= piov->to[*iov_index].iov_len; + (*iov_index)++; + } + if (*iov_index >= piov->nr) { + pr_err("Async decompression ran out of destination iovecs\n"); + return -1; + } + + dst = (char *)piov->to[*iov_index].iov_base + *iov_offset; + chunk = piov->to[*iov_index].iov_len - *iov_offset; + if (chunk > bytes) + chunk = bytes; + + if (src) { + memcpy(dst, src, chunk); + src += chunk; + } else { + memset(dst, 0, chunk); + } + *iov_offset += chunk; + bytes -= chunk; + } + + return 0; +} + +static void encoded_read_ctx_fini(struct encoded_read_ctx *ctx) +{ + if (!ctx) + return; + encoded_read_ctx_end_work(ctx); + decompression_pool_destroy(ctx->pool); +} + +static bool page_read_chain_has_encoded_async(struct page_read *pr) +{ + struct page_read_iov *piov; + + list_for_each_entry(piov, &pr->async, l) { + if (piov->storage == VMA_IO_ENCODED) + return true; + } + + return pr->parent && page_read_chain_has_encoded_async(pr->parent); +} + +/* + * Restore one bounded encoded piov. Validate its immutable block metadata, + * read the packed payload once, complete raw blocks inline, and describe + * zero/LZ4 blocks as independent worker jobs. Blocks spanning multiple + * destination iovecs use the reusable scratch buffer. The caller owns queue + * removal; this routine leaves the piov intact on both success and failure. + */ +static int process_encoded_async_read(int fd, struct page_read_iov *piov, + struct encoded_read_ctx *ctx, + bool payload_ready, + struct encoded_prefetch *prefetch) +{ + struct decompress_job *jobs; + char *compressed; + char *scratch; + size_t scratch_cap; + size_t compressed_offset = 0; + size_t total_compressed = piov->total_compressed_size; + size_t jobs_uncompressed = 0; + size_t output_bytes = 0; + size_t expected_output; + size_t nr_jobs = 0; + unsigned int iov_index = 0; + size_t iov_offset = 0; + bool parallel_zero = false; + size_t i; int ret = -1; - (void)pr; - - /* - * An all-zero batch has total == 0: there is nothing - * to read from the image and every block is rebuilt by - * memset below, so leave comp_buf NULL. xmalloc(0) may - * return NULL, which must not be treated as an error. - * Route a genuine allocation failure through the err - * path so the async queue is drained. - */ - comp_buf = NULL; - if (total) { - comp_buf = xmalloc(total); - if (!comp_buf) - goto out; - - if (pread_full(fd, comp_buf, total, piov->from)) - goto out; + /* Validate counts before they drive buffer growth or offset arithmetic. */ + if (!ctx || !ctx->batch_acquired) { + pr_err("Encoded async I/O job has no active shared read context\n"); + return -1; } - - if (piov->region_pages) { - region_scratch_cap = (size_t)piov->region_pages * PAGE_SIZE; - region_scratch = xmalloc(region_scratch_cap); - if (!region_scratch) - goto out; + parallel_zero = compressed_restore_has_parallel_capacity(opts.decompress_threads); + if (!piov->n_compressed_size || !piov->compressed_size) { + pr_err("Encoded async I/O job has no block metadata\n"); + return -1; } + if (piov->n_compressed_size > (size_t)INT_MAX || + piov->n_compressed_size > SIZE_MAX / sizeof(*jobs)) { + pr_err("Encoded async I/O job has too many blocks: %zu\n", + piov->n_compressed_size); + return -1; + } + if (piov->region_pages && !piov->block_pages) { + pr_err("Encoded region I/O job has no block-page metadata\n"); + return -1; + } + if ((uint64_t)total_compressed != piov->total_compressed_size) { + pr_err("Encoded async I/O size does not fit in size_t\n"); + return -1; + } + if (piov->n_pages > SIZE_MAX / PAGE_SIZE) { + pr_err("Encoded async I/O page count does not fit in size_t\n"); + return -1; + } + expected_output = (size_t)piov->n_pages * (size_t)PAGE_SIZE; + /* Grow buffers under the active read's bounded working-set lease. */ - comp_off = 0; - for (int i = 0; i < piov->n_compressed_size; i++) { - uint32_t cs = piov->compressed_size[i]; - unsigned int block_pages; - size_t block_bytes; - size_t bound; - size_t out_left; - char *src; + if (ctx->jobs_cap < piov->n_compressed_size) { + void *new_jobs = xrealloc(ctx->jobs, + piov->n_compressed_size * sizeof(*jobs)); - if (piov->region_pages) { - block_pages = piov->block_pages[i]; - bound = REGION_COMPRESSED_SIZE_BOUND(block_pages); - } else { - block_pages = 1; - bound = PAGE_COMPRESSED_SIZE_BOUND; - } - block_bytes = (size_t)block_pages * PAGE_SIZE; - - if (cs > bound) { - pr_err("Async: invalid compressed size %u for block %d\n", - cs, i); + if (!new_jobs) goto out; - } + ctx->jobs = new_jobs; + ctx->jobs_cap = piov->n_compressed_size; + } + jobs = ctx->jobs; - /* - * Decompress this block to a contiguous - * source buffer. - * - * Region mode raw blocks (cs == block_bytes) - * skip the staging memcpy and copy directly - * out of comp_buf. Zero blocks (cs == 0) are - * handled by memset on the destination. - * - * Compressed blocks normally decompress into - * region_scratch and we then memcpy into the - * iovecs. As an optimisation, when the whole - * block fits inside a single destination iovec - * (the common case after a contiguous VMA), we - * decompress directly into the iovec and skip - * the staging buffer. - */ - if (piov->region_pages) { - if (cs == 0) { - src = NULL; /* memset on dst */ - } else if (cs == block_bytes) { - src = comp_buf + comp_off; - } else { - /* Locate destination iovec for this block */ - int peek_iov = iov_idx; - size_t peek_off = iov_off; + if (total_compressed) { + if (ctx->compressed_cap < total_compressed) { + void *new_compressed; - while (peek_iov < piov->nr && peek_off >= piov->to[peek_iov].iov_len) { - peek_off -= piov->to[peek_iov].iov_len; - peek_iov++; - } - - if (peek_iov < piov->nr && - peek_off + block_bytes <= - piov->to[peek_iov].iov_len) { - /* Block fits — decompress directly into iovec */ - char *direct = (char *)piov->to[peek_iov].iov_base + peek_off; - - if (decompress_region(comp_buf + comp_off, - cs, block_pages, - direct)) { - pr_err("Async region decompress failed at block %d\n", - i); - goto out; - } - /* Advance iovec cursor and skip the copy loop. */ - iov_off += block_bytes; - comp_off += cs; - continue; - } - - if (decompress_region(comp_buf + comp_off, - cs, block_pages, - region_scratch)) { - pr_err("Async region decompress failed at block %d\n", - i); - goto out; - } - src = region_scratch; - } - } else { - src = NULL; /* handled below */ - } - - /* - * Walk the destination iovec array, - * copying block_bytes worth of decompressed - * data, splitting across iovecs as needed. - */ - out_left = block_bytes; - while (out_left > 0) { - char *dst; - size_t this_chunk; - - while (iov_idx < piov->nr && iov_off >= piov->to[iov_idx].iov_len) { - iov_off -= piov->to[iov_idx].iov_len; - iov_idx++; - } - if (iov_idx >= piov->nr) { - pr_err("Async: ran out of iovs at block %d\n", i); + if (payload_ready) { + pr_err("Prefetched encoded payload exceeds its buffer\n"); goto out; } + new_compressed = xrealloc(ctx->compressed, + total_compressed); - dst = (char *)piov->to[iov_idx].iov_base + iov_off; - this_chunk = piov->to[iov_idx].iov_len - iov_off; - if (this_chunk > out_left) - this_chunk = out_left; + if (!new_compressed) + goto out; + ctx->compressed = new_compressed; + ctx->compressed_cap = total_compressed; + } + compressed = ctx->compressed; + if (!payload_ready && + pread_full(fd, compressed, total_compressed, piov->from)) + goto out; + } else + compressed = NULL; + scratch = ctx->scratch; + scratch_cap = ctx->scratch_cap; - if (piov->region_pages) { - if (src) { - memcpy(dst, src, this_chunk); - src += this_chunk; - } else { - /* zero block */ - memset(dst, 0, this_chunk); - } - } else { - /* Per-page blocks fit in one page-aligned iovec slice. */ - if (cs == 0) - memset(dst, 0, this_chunk); - else if (cs == PAGE_SIZE) - memcpy(dst, comp_buf + comp_off, this_chunk); - else if (decompress_data(comp_buf + comp_off, cs, PAGE_SIZE, dst)) { - pr_err("Async decompress failed for page %d\n", i); - goto out; - } + /* Build disjoint zero/LZ4 jobs while completing raw blocks inline. */ + for (i = 0; i < piov->n_compressed_size; i++) { + uint32_t compressed_size = piov->compressed_size[i]; + unsigned int block_pages = 1; + size_t block_bytes; + size_t bound = PAGE_COMPRESSED_SIZE_BOUND; + char *direct_dst = NULL; + + if (piov->region_pages) + block_pages = piov->block_pages[i]; + if (!block_pages || + (piov->region_pages && block_pages > piov->region_pages)) { + pr_err("Async: invalid page count %u for block %zu\n", + block_pages, i); + goto out; + } + block_bytes = (size_t)block_pages * PAGE_SIZE; + if (piov->region_pages) + bound = REGION_COMPRESSED_SIZE_BOUND(block_pages); + if (compressed_size > bound || + compressed_size > total_compressed - compressed_offset) { + pr_err("Async: invalid compressed size %u for block %zu\n", + compressed_size, i); + goto out; + } + if (output_bytes > SIZE_MAX - block_bytes) { + pr_err("Async decompressed size overflows\n"); + goto out; + } + output_bytes += block_bytes; + + while (iov_index < piov->nr && iov_offset >= piov->to[iov_index].iov_len) { + iov_offset -= piov->to[iov_index].iov_len; + iov_index++; + } + if (iov_index >= piov->nr) { + pr_err("Async: ran out of iovecs at block %zu\n", i); + goto out; + } + if (block_bytes <= piov->to[iov_index].iov_len - iov_offset) + direct_dst = (char *)piov->to[iov_index].iov_base + iov_offset; + + if (!compressed_size && (!direct_dst || !parallel_zero)) { + if (transfer_async_block(piov, &iov_index, &iov_offset, + NULL, block_bytes)) + goto out; + } else if (compressed_size == block_bytes) { + if (transfer_async_block(piov, &iov_index, &iov_offset, + compressed + compressed_offset, + block_bytes)) + goto out; + } else if (compressed_size >= block_bytes) { + pr_err("Async: LZ4 block %zu has invalid size %u for %zu bytes\n", + i, compressed_size, block_bytes); + goto out; + } else if (direct_dst) { + struct decompress_job *job = &jobs[nr_jobs++]; + + job->src = compressed_size ? compressed + compressed_offset : NULL; + job->dst = direct_dst; + job->compressed_size = compressed_size; + job->pages = block_pages; + job->block_index = i; + if (jobs_uncompressed > SIZE_MAX - block_bytes) { + pr_err("Parallel decompression size overflows\n"); + goto out; } + jobs_uncompressed += block_bytes; + iov_offset += block_bytes; + } else { + void *new_scratch; - iov_off += this_chunk; - out_left -= this_chunk; + /* A region crossing iovecs needs one serial staging copy. */ + if (scratch_cap < block_bytes) { + new_scratch = xrealloc(scratch, block_bytes); + if (!new_scratch) + goto out; + scratch = new_scratch; + scratch_cap = block_bytes; + ctx->scratch = scratch; + ctx->scratch_cap = scratch_cap; + } + if (decompress_region(compressed + compressed_offset, + compressed_size, block_pages, + scratch)) { + pr_err("Async decompression failed at split block %zu\n", i); + goto out; + } + if (transfer_async_block(piov, &iov_index, &iov_offset, + scratch, block_bytes)) + goto out; } - comp_off += cs; + compressed_offset += compressed_size; } + while (iov_index < piov->nr && iov_offset >= piov->to[iov_index].iov_len) { + iov_offset -= piov->to[iov_index].iov_len; + iov_index++; + } + if (compressed_offset != total_compressed || iov_index != piov->nr || + iov_offset || output_bytes != expected_output) { + pr_err("Inconsistent encoded async I/O sizes: payload=%zu/%zu output=%zu/%zu\n", + compressed_offset, total_compressed, output_bytes, + expected_output); + goto out; + } + + /* The caller can read the next payload while pool workers run these jobs. */ + if (decompress_jobs_parallel_pool_with_caller_work( + &ctx->pool, jobs, nr_jobs, jobs_uncompressed, + opts.decompress_threads, + prefetch ? encoded_prefetch_read : NULL, prefetch)) + goto out; + ret = 0; out: - xfree(comp_buf); - xfree(region_scratch); return ret; } -static int process_async_reads(struct page_read *pr) +static bool encoded_prefetch_eligible(const struct page_read *pr, + const struct page_read_iov *current, + const struct page_read_iov *next) +{ + size_t next_size = next->total_compressed_size; + + if (opts.stream || pr->use_direct || next->storage != VMA_IO_ENCODED) + return false; + if (!compressed_restore_has_parallel_capacity(opts.decompress_threads)) + return false; + if (current->n_pages < PARALLEL_RESTORE_MIN_BATCH_BYTES / PAGE_SIZE || + next->n_pages < PARALLEL_RESTORE_MIN_BATCH_BYTES / PAGE_SIZE) + return false; + if (!next_size || next_size > ASYNC_BATCH_MAX_BYTES || + (uint64_t)next_size != next->total_compressed_size || + next->end < next->from) + return false; + if ((uint64_t)(next->end - next->from) != + next->total_compressed_size) + return false; + + return true; +} + +/* + * Drain one page-reader chain in image order. Each queue element has one of + * four storage kinds: zero and packed raw bypass LZ4, ordinary entries use + * preadv(), and encoded entries use the shared chain context above. Parent + * queues share the same context. Acquire its working-set lease only when an + * encoded request is reached. A second nonblocking lease permits one payload + * read to overlap decoding. Release both before ordinary I/O or a parent + * queue. Any error drains the remaining subtree before close_page_read() + * checks that all queues are empty. + */ +static int process_async_reads_ctx(struct page_read *pr, + struct encoded_read_ctx *encoded_ctx) { int fd, ret = 0; struct page_read_iov *piov, *n; off_t first_off = OFF_MAX, last_end = OFF_MIN; fd = img_raw_fd(pr->pi); + /* Hint bounded nearby ranges before issuing their explicit preadv calls. */ if (!pr->use_direct) { list_for_each_entry(piov, &pr->async, l) { bool merge = first_off != OFF_MAX && @@ -1819,6 +2548,7 @@ static int process_async_reads(struct page_read *pr) } } + /* Consume and free each request only after its destination is complete. */ list_for_each_entry_safe(piov, n, &pr->async, l) { ssize_t ret; struct iovec *iovs = piov->to; @@ -1826,14 +2556,74 @@ static int process_async_reads(struct page_read *pr) pr_debug("Read piov iovs %d, from %ju, len %ju, first %p:%zu\n", piov->nr, piov->from, piov->end - piov->from, piov->to->iov_base, piov->to->iov_len); - if (piov->n_compressed_size) { - if (process_compressed_async_read(pr, fd, piov)) { + /* Do not occupy an encoded-work slot while serving ordinary ranges. */ + if (piov->storage != VMA_IO_ENCODED) + encoded_read_ctx_end_work(encoded_ctx); + + if (piov->storage == VMA_IO_ZERO) { + if (validate_direct_compressed_iov(piov)) { ret = -1; goto err; } + for (unsigned int i = 0; i < piov->nr; i++) + memset(piov->to[i].iov_base, 0, piov->to[i].iov_len); goto next; } + if (piov->storage == VMA_IO_PACKED_RAW) { + if (validate_direct_compressed_iov(piov)) { + ret = -1; + goto err; + } + goto more; + } + + if (piov->storage == VMA_IO_ENCODED) { + bool prefetch_prepared = false; + int payload_ready; + + if (!encoded_ctx) { + pr_err("Encoded async I/O job has no shared read context\n"); + ret = -1; + goto err; + } + if (!encoded_ctx->batch_acquired) + encoded_read_ctx_begin_work(encoded_ctx); + + payload_ready = encoded_prefetch_take(encoded_ctx, piov); + if (payload_ready < 0) { + ret = -1; + goto err; + } + if (!list_is_last(&piov->l, &pr->async) && + encoded_prefetch_eligible(pr, piov, n)) { + prefetch_prepared = encoded_prefetch_prepare( + encoded_ctx, fd, n->from, + (size_t)n->total_compressed_size); + } else if (encoded_ctx->prefetch_batch_acquired) { + encoded_prefetch_disable(encoded_ctx); + } + + ret = process_encoded_async_read( + fd, piov, encoded_ctx, payload_ready > 0, + prefetch_prepared ? &encoded_ctx->prefetch : NULL); + if (prefetch_prepared) { + if (ret < 0 || !encoded_ctx->prefetch.complete) + encoded_prefetch_disable(encoded_ctx); + else + encoded_prefetch_publish(encoded_ctx, n); + } + if (ret < 0) + goto err; + goto next; + } + if (piov->storage != VMA_IO_UNCOMPRESSED) { + pr_err("Unknown async I/O storage kind %d\n", + piov->storage); + ret = -1; + goto err; + } + more: ret = preadv(fd, piov->to, piov->nr, piov->from); if (fault_injected(FI_PARTIAL_PAGES)) { @@ -1861,7 +2651,8 @@ static int process_async_reads(struct page_read *pr) goto err; } - if (opts.auto_dedup && punch_hole(pr, piov->from, ret, false)) + if (opts.auto_dedup && piov->storage != VMA_IO_PACKED_RAW && + punch_hole(pr, piov->from, ret, false)) goto err; if (ret != piov->end - piov->from) { @@ -1886,9 +2677,25 @@ static int process_async_reads(struct page_read *pr) xfree(iovs); xfree(piov); } + /* Parent readahead and raw prefixes must not inherit an idle lease. */ + encoded_read_ctx_end_work(encoded_ctx); + if (pr->parent) { + ret = process_async_reads_ctx(pr->parent, encoded_ctx); + if (ret) + return ret; + } - if (pr->parent) - ret = process_async_reads(pr->parent); + /* + * A final auto-dedup batch used to be deferred until close(), whose + * void callback cannot report a failed fallocate(). Flush it while the + * caller can still propagate the error from ->sync(). + */ + if (pr->bunch.iov_len > 0) { + ret = punch_hole(pr, 0, 0, true); + if (ret) + return ret; + pr->bunch.iov_len = 0; + } return ret; err: @@ -1896,17 +2703,53 @@ err: return -1; } -static void close_page_read(struct page_read *pr) +static int process_async_reads(struct page_read *pr) { + struct page_read *owner = pr->encoded_read_owner; + struct encoded_read_ctx *ctx = NULL; int ret; + if (!owner) { + pr_err("Page reader has no encoded-read context owner\n"); + drain_async_queue(pr); + return -1; + } + + /* Raw/zero/uncompressed-only syncs need neither buffers nor a lease. */ + if (page_read_chain_has_encoded_async(pr)) { + if (!owner->encoded_read_ctx) { + owner->encoded_read_ctx = + xzalloc(sizeof(*owner->encoded_read_ctx)); + } + if (!owner->encoded_read_ctx) { + drain_async_queue(pr); + return -1; + } + ctx = owner->encoded_read_ctx; + } + + ret = process_async_reads_ctx(pr, ctx); + encoded_read_ctx_end_work(ctx); + return ret; +} + +static void close_page_read(struct page_read *pr) +{ BUG_ON(!list_empty(&pr->async)); + /* + * Restore tasks close their page readers before they fork children or + * remap the PIE bootstrap. Page-server readers close after their last + * sync, so one pool also spans its bounded decode chunks. + */ + if (pr->encoded_read_owner == pr) { + encoded_read_ctx_fini(pr->encoded_read_ctx); + xfree(pr->encoded_read_ctx); + pr->encoded_read_ctx = NULL; + } if (pr->bunch.iov_len > 0) { - ret = punch_hole(pr, 0, 0, true); - if (ret == -1) - return; - + /* punch_hole() logs failures; cleanup must run in either case. */ + (void)punch_hole(pr, 0, 0, true); pr->bunch.iov_len = 0; } @@ -1930,6 +2773,7 @@ static void reset_pagemap(struct page_read *pr) { pr->cvaddr = 0; pr->pi_off = 0; + pr->stream_padding = 0; pr->compressed_size_index = 0; pr->region_block_offset = 0; pr->curr_pme = -1; @@ -1941,6 +2785,11 @@ static void reset_pagemap(struct page_read *pr) reset_pagemap(pr->parent); } +/* + * Open one optional parent reader. Until the final assignment to pr->parent, + * this function owns both the parent-directory fd and the allocated reader; + * each failure label releases exactly the resources acquired above it. + */ static int try_open_parent(int dfd, unsigned long id, struct page_read *pr, int pr_flags) { int pfd, ret; @@ -1981,6 +2830,14 @@ err: return -1; } +static void set_encoded_read_owner(struct page_read *pr, + struct page_read *owner) +{ + pr->encoded_read_owner = owner; + if (pr->parent) + set_encoded_read_owner(pr->parent, owner); +} + static void init_compat_pagemap_entry(PagemapEntry *pe) { /* @@ -2012,8 +2869,14 @@ static int validate_compressed_pagemap_entry(PagemapEntry *pe) if (pe->has_region_pages && pe->region_pages) region_pages = pe->region_pages; - if (!pe->n_compressed_size) + if (!pe->n_compressed_size) { + if (pe->has_total_compressed_size || pe->has_region_pages) { + pr_err("Compression metadata without block sizes on pagemap entry %#" PRIx64 "\n", + pe->vaddr); + return -1; + } return 0; + } #ifndef CONFIG_LZ4 pr_err("Pagemap contains compressed pages but CRIU was built without LZ4 support (CONFIG_LZ4)\n"); @@ -2026,15 +2889,11 @@ static int validate_compressed_pagemap_entry(PagemapEntry *pe) return -1; } - if (!pe->nr_pages) { - pr_err("Compressed pagemap entry %#" PRIx64 " has no pages\n", pe->vaddr); - return -1; - } - if (region_pages) { if (region_pages > MAX_REGION_PAGES) { - pr_err("Compressed pagemap entry %#" PRIx64 " has invalid region_pages %u\n", - pe->vaddr, region_pages); + pr_err("Compressed pagemap entry %#" PRIx64 + " has invalid region_pages %u (max %lu)\n", + pe->vaddr, region_pages, MAX_REGION_PAGES); return -1; } expected_blocks = pe->nr_pages / region_pages; @@ -2097,6 +2956,92 @@ static int validate_compressed_pagemap_entry(PagemapEntry *pe) return 0; } +static int validate_pagemap_entry_layout(PagemapEntry *pe, + uint64_t *previous_end) +{ + uint64_t length; + uint64_t end; + + if (pagemap_payload_aligned(pe) && !pagemap_present(pe)) { + pr_err("Aligned payload flag on non-present pagemap entry %#" PRIx64 + "\n", pe->vaddr); + return -1; + } + if (pagemap_present(pe) && pagemap_in_parent(pe)) { + pr_err("Pagemap entry %#" PRIx64 + " cannot be present and inherited at the same time\n", + pe->vaddr); + return -1; + } + + if (pe->vaddr % PAGE_SIZE) { + pr_err("Pagemap entry address %#" PRIx64 " is not page-aligned\n", + pe->vaddr); + return -1; + } + if (!pe->nr_pages) { + pr_err("Pagemap entry %#" PRIx64 " has no pages\n", pe->vaddr); + return -1; + } + if (pe->nr_pages > UINT64_MAX / PAGE_SIZE) { + pr_err("Pagemap entry %#" PRIx64 " page count overflows\n", + pe->vaddr); + return -1; + } + length = pe->nr_pages * PAGE_SIZE; + if (pe->vaddr > UINT64_MAX - length) { + pr_err("Pagemap entry %#" PRIx64 " end overflows\n", pe->vaddr); + return -1; + } + end = pe->vaddr + length; + if (pe->vaddr > ULONG_MAX || end > ULONG_MAX) { + pr_err("Pagemap entry %#" PRIx64 "-%#" PRIx64 + " does not fit in an address\n", pe->vaddr, end); + return -1; + } + if (pe->vaddr < *previous_end) { + pr_err("Pagemap entry %#" PRIx64 "-%#" PRIx64 + " overlaps or precedes the previous entry ending at %#" PRIx64 "\n", + pe->vaddr, end, *previous_end); + return -1; + } + + *previous_end = end; + return 0; +} + +static int validate_pages_image_layout(PagemapEntry *pe, off_t *offset) +{ + uint64_t payload = 0; + size_t i; + + if (!pagemap_present(pe)) + return 0; + + if (pagemap_payload_aligned(pe)) { + if (*offset < 0 || *offset > OFF_MAX - (PAGE_SIZE - 1)) { + pr_err("Pages image offset %jd cannot be page-aligned for entry %#" PRIx64 "\n", + (intmax_t)*offset, pe->vaddr); + return -1; + } + *offset = pagemap_page_align_offset(*offset); + } + + if (pe->n_compressed_size) { + for (i = 0; i < pe->n_compressed_size; i++) + payload += pe->compressed_size[i]; + } else { + payload = pe->nr_pages * PAGE_SIZE; + } + if (payload > (uint64_t)(OFF_MAX - *offset)) { + pr_err("Pages image payload for entry %#" PRIx64 + " exceeds the representable file offset\n", pe->vaddr); + return -1; + } + *offset += (off_t)payload; + return 0; +} + static bool page_read_has_compressed_entries(struct page_read *pr) { int i; @@ -2110,38 +3055,229 @@ static bool page_read_has_compressed_entries(struct page_read *pr) } /* - * The pagemap entry size is at least 8 bytes for small mappings with - * low address and may get to 18 bytes or even more for large mappings - * with high address and in_parent flag set. 16 seems to be nice round - * number to minimize {over,under}-allocations + * Inspect [start, end) without advancing any page-reader cursor. Return 1 + * when an overlapping block requires LZ4 decoding. With @premap_mixed, also + * request premapping when splitting raw and zero blocks would exceed the + * bounded number of direct restorer jobs, or when a large zero run can use + * parallel filling. Parent entries are inspected in their owning reader. + * Return 0 when PIE can restore the range directly and -1 for invalid metadata. */ -#define PAGEMAP_ENTRY_SIZE_ESTIMATE 16 +static int page_read_range_needs_decode(struct page_read *pr, + unsigned long start, + unsigned long end, + bool premap_mixed) +{ + unsigned int direct_runs = 0; + unsigned int previous_region_pages = 0; + enum restore_vma_io_storage previous_storage = VMA_IO_UNCOMPRESSED; + bool have_previous_storage = false; + bool parallel_zero = false; + int left = 0; + int right; + int i; + + if (!pr || start >= end || start % PAGE_SIZE || end % PAGE_SIZE) { + pr_err("Invalid compressed-page range %#lx-%#lx\n", start, end); + return -1; + } + if (premap_mixed) + parallel_zero = compressed_restore_has_parallel_capacity(opts.decompress_threads); + + /* Find the first pagemap entry whose end is after start. */ + right = pr->nr_pmes; + while (left < right) { + int middle = left + (right - left) / 2; + PagemapEntry *pe = pr->pmes[middle]; + unsigned long pe_start = + (unsigned long)decode_pointer(pe->vaddr); + unsigned long pe_end = pe_start + pagemap_len(pe); + + if (pe_end <= start) + left = middle + 1; + else + right = middle; + } + + /* Classify only the compressed blocks that overlap the requested range. */ + for (i = left; i < pr->nr_pmes; i++) { + PagemapEntry *pe = pr->pmes[i]; + unsigned long pe_start = + (unsigned long)decode_pointer(pe->vaddr); + unsigned long pe_end = pe_start + pagemap_len(pe); + unsigned long overlap_start; + unsigned long overlap_end; + unsigned int region_pages; + unsigned long zero_bytes = 0; + unsigned int zero_blocks = 0; + bool parallel_zero_entry = parallel_zero; + size_t first_block; + size_t last_block; + size_t block; + + if (pe_start >= end) + break; + if (pe_end <= start) + continue; + + overlap_start = max(start, pe_start); + overlap_end = min(end, pe_end); + if (pagemap_in_parent(pe)) { + int ret; + + if (!pr->parent) { + pr_err("Pagemap range %#lx-%#lx has no parent reader\n", + overlap_start, overlap_end); + return -1; + } + ret = page_read_range_needs_decode(pr->parent, + overlap_start, + overlap_end, + premap_mixed); + if (ret) + return ret; + have_previous_storage = false; + continue; + } + if (!pe->n_compressed_size) { + have_previous_storage = false; + continue; + } + + region_pages = pagemap_region_pages(pe); + if (have_previous_storage && + region_pages != previous_region_pages) + have_previous_storage = false; + if (region_pages) { + uint64_t first_page = + (overlap_start - pe_start) / PAGE_SIZE; + uint64_t end_page = + (overlap_end - pe_start) / PAGE_SIZE; + + if (first_page % region_pages || + (end_page % region_pages && overlap_end != pe_end)) + parallel_zero_entry = false; + first_block = first_page / region_pages; + last_block = (end_page + region_pages - 1) / + region_pages; + } else { + first_block = (overlap_start - pe_start) / PAGE_SIZE; + last_block = (overlap_end - pe_start) / PAGE_SIZE; + } + + if (last_block > pe->n_compressed_size) { + pr_err("LZ4 range block index %zu exceeds %zu blocks\n", + last_block, pe->n_compressed_size); + return -1; + } + for (block = first_block; block < last_block; block++) { + enum restore_vma_io_storage storage; + unsigned int block_pages; + + if (compressed_block_storage(pe, block, &storage, + &block_pages)) + return -1; + if (storage == VMA_IO_ENCODED) + return 1; + if (parallel_zero_entry && storage == VMA_IO_ZERO) { + uint64_t first_page = block; + unsigned long block_start; + unsigned long block_end; + + if (region_pages) + first_page *= region_pages; + block_start = pe_start + first_page * PAGE_SIZE; + block_end = block_start + (unsigned long)block_pages * PAGE_SIZE; + + if (block_start < overlap_start || block_end > overlap_end) { + zero_bytes = 0; + zero_blocks = 0; + } else { + zero_bytes += block_end - block_start; + zero_blocks++; + if (zero_blocks > 1 && + zero_bytes >= PARALLEL_RESTORE_MIN_BATCH_BYTES) + return 1; + } + } else { + zero_bytes = 0; + zero_blocks = 0; + } + /* Each storage transition creates another direct restore job. */ + if (premap_mixed && + (!have_previous_storage || storage != previous_storage) && + ++direct_runs > DIRECT_COMPRESSED_RUN_MAX) + return 1; + previous_storage = storage; + previous_region_pages = region_pages; + have_previous_storage = true; + } + } + + return 0; +} + +int page_read_range_has_lz4(struct page_read *pr, unsigned long start, + unsigned long end) +{ + return page_read_range_needs_decode(pr, start, end, false); +} + +int page_read_range_needs_premap(struct page_read *pr, unsigned long start, + unsigned long end) +{ + return page_read_range_needs_decode(pr, start, end, true); +} + +int page_read_range_has_parent(struct page_read *pr, unsigned long start, + unsigned long end) +{ + int left = 0; + int right; + int i; + + if (!pr || start >= end || start % PAGE_SIZE || end % PAGE_SIZE) { + pr_err("Invalid parent-page range %#lx-%#lx\n", start, end); + return -1; + } + + right = pr->nr_pmes; + while (left < right) { + int middle = left + (right - left) / 2; + PagemapEntry *pe = pr->pmes[middle]; + unsigned long pe_start = + (unsigned long)decode_pointer(pe->vaddr); + unsigned long pe_end = pe_start + pagemap_len(pe); + + if (pe_end <= start) + left = middle + 1; + else + right = middle; + } + + for (i = left; i < pr->nr_pmes; i++) { + PagemapEntry *pe = pr->pmes[i]; + unsigned long pe_start = + (unsigned long)decode_pointer(pe->vaddr); + unsigned long pe_end = pe_start + pagemap_len(pe); + + if (pe_start >= end) + break; + if (pe_end > start && pagemap_in_parent(pe)) + return 1; + } + + return 0; +} + +#define PAGEMAP_INITIAL_ENTRIES 64 static int init_pagemaps(struct page_read *pr) { - off_t fsize; - int nr_pmes, nr_realloc; + uint64_t previous_end = 0; + off_t pages_offset = 0; + size_t capacity = PAGEMAP_INITIAL_ENTRIES; - if (opts.stream) { - /* - * TODO - There is no easy way to estimate the size of the - * pagemap that is still to be read from the pipe. Possible - * solution is to ask the image streamer for the size of the - * image. 1024 is a wild guess (more space is allocated if - * needed). - */ - fsize = 1024; - } else { - fsize = img_raw_size(pr->pmi); - } - - if (fsize < 0) - return -1; - - nr_pmes = fsize / PAGEMAP_ENTRY_SIZE_ESTIMATE + 1; - nr_realloc = nr_pmes / 2; - - pr->pmes = xzalloc(nr_pmes * sizeof(*pr->pmes)); + pr->pmes = xzalloc(capacity * sizeof(*pr->pmes)); if (!pr->pmes) return -1; @@ -2149,9 +3285,35 @@ static int init_pagemaps(struct page_read *pr) pr->curr_pme = -1; while (1) { + PagemapEntry **new; PagemapEntry *pe; - int ret = pb_read_one_eof(pr->pmi, &pr->pmes[pr->nr_pmes], PB_PAGEMAP); + int ret; + if ((size_t)pr->nr_pmes == capacity) { + size_t alloc_size; + size_t new_capacity; + + if (capacity >= INT_MAX || + __builtin_mul_overflow(capacity, (size_t)2, + &new_capacity)) { + pr_err("Too many pagemap entries\n"); + goto free_pagemaps; + } + new_capacity = min(new_capacity, (size_t)INT_MAX); + if (__builtin_mul_overflow(new_capacity, sizeof(*pr->pmes), + &alloc_size)) { + pr_err("Pagemap entry array size overflows\n"); + goto free_pagemaps; + } + new = xrealloc(pr->pmes, alloc_size); + if (!new) + goto free_pagemaps; + pr->pmes = new; + capacity = new_capacity; + } + + ret = pb_read_one_eof(pr->pmi, &pr->pmes[pr->nr_pmes], + PB_PAGEMAP); if (ret < 0) goto free_pagemaps; if (ret == 0) @@ -2159,17 +3321,12 @@ static int init_pagemaps(struct page_read *pr) pe = pr->pmes[pr->nr_pmes++]; init_compat_pagemap_entry(pe); + if (validate_pagemap_entry_layout(pe, &previous_end)) + goto free_pagemaps; if (validate_compressed_pagemap_entry(pe)) goto free_pagemaps; - - if (pr->nr_pmes >= nr_pmes) { - PagemapEntry **new; - nr_pmes += nr_realloc; - new = xrealloc(pr->pmes, nr_pmes * sizeof(*pr->pmes)); - if (!new) - goto free_pagemaps; - pr->pmes = new; - } + if (validate_pages_image_layout(pe, &pages_offset)) + goto free_pagemaps; } close_image(pr->pmi); @@ -2272,11 +3429,14 @@ int open_page_read_at(int dfd, unsigned long img_id, struct page_read *pr, int p pr->parent = NULL; pr->cvaddr = 0; pr->pi_off = 0; + pr->stream_padding = 0; pr->compressed_size_index = 0; pr->region_block_offset = 0; pr->cached_region = NULL; pr->cached_region_vaddr = 0; pr->cached_region_size = 0; + pr->encoded_read_ctx = NULL; + pr->encoded_read_owner = pr; pr->bunch.iov_len = 0; pr->bunch.iov_base = NULL; pr->pmes = NULL; @@ -2297,6 +3457,7 @@ int open_page_read_at(int dfd, unsigned long img_id, struct page_read *pr, int p close_image(pr->pmi); return -1; } + set_encoded_read_owner(pr, pr); pr->pi = open_pages_image_at(dfd, flags, pr->pmi, &pr->pages_img_id); if (!pr->pi) { @@ -2360,8 +3521,7 @@ int open_page_read_at(int dfd, unsigned long img_id, struct page_read *pr, int p pr->maybe_read_page = maybe_read_page_img_streamer; else { pr->maybe_read_page = maybe_read_page_local_compressed; - if (!pr->parent && !opts.lazy_pages && - !page_read_has_compressed_entries(pr)) + if (!pr->parent && !opts.lazy_pages) pr->pieok = true; } @@ -2396,5 +3556,12 @@ void dup_page_read(struct page_read *src, struct page_read *dst) dst->cached_region = NULL; dst->cached_region_vaddr = 0; dst->cached_region_size = 0; + dst->encoded_read_ctx = NULL; + /* + * UFFD fork readers are shallow duplicates and keep their root lpi alive + * through its reference count. Reuse that root's chain context instead of + * allocating an unowned context that lpi_fini() cannot release. + */ + dst->encoded_read_owner = src->encoded_read_owner; dst->reset(dst); } diff --git a/criu/pie/restorer.c b/criu/pie/restorer.c index 553719c43..59b7c6b67 100644 --- a/criu/pie/restorer.c +++ b/criu/pie/restorer.c @@ -1,5 +1,6 @@ #include #include +#include #include #include @@ -1680,8 +1681,10 @@ static int process_aio_event(struct task_restore_args *args, aio_context_t aio_c } cb = (struct iocb *)(unsigned long)ev->obj; + idx = cb - iocbs; + r = rio_ptrs[idx]; - if (args->auto_dedup) { + if (args->auto_dedup && r->storage == VMA_IO_UNCOMPRESSED) { long fr = sys_fallocate(fd, FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE, cb->aio_offset, res); if (fr < 0) @@ -1691,8 +1694,6 @@ static int process_aio_event(struct task_restore_args *args, aio_context_t aio_c if (res == ev->data) return 0; - idx = cb - iocbs; - r = rio_ptrs[idx]; advance_vma_io_retry(r, res, &iov_next, &nr_next); if (!iov_next || nr_next == 0) { pr_err("AIO retry advance produced no work after %zd bytes\n", res); @@ -1760,6 +1761,148 @@ static int reap_aio_events(struct task_restore_args *args, aio_context_t aio_ctx return 0; } +static int restore_vma_preadv_one(struct task_restore_args *args, + struct restore_vma_io *rio, + bool allow_dedup); + +static int validate_direct_vma_io(struct restore_vma_io *rio) +{ + uint64_t payload_bytes = 0; + uint64_t output_bytes = 0; + uint64_t iov_bytes = 0; + int i; + + if (rio->storage != VMA_IO_PACKED_RAW && rio->storage != VMA_IO_ZERO) + return -1; + if (rio->n_compressed_size <= 0 || !rio->compressed_size) { + pr_err("Direct compressed VMA IO has no block metadata\n"); + return -1; + } + if (rio->region_pages && !rio->block_pages) { + pr_err("Direct compressed region VMA IO has no page counts\n"); + return -1; + } + + for (i = 0; i < rio->n_compressed_size; i++) { + unsigned int block_pages = rio->region_pages ? + rio->block_pages[i] : 1; + uint64_t block_bytes; + uint32_t compressed_size = rio->compressed_size[i]; + + if (!block_pages || + (rio->region_pages && block_pages > rio->region_pages)) { + pr_err("Invalid direct VMA IO block page count %u\n", + block_pages); + return -1; + } + block_bytes = (uint64_t)block_pages * PAGE_SIZE; + if (rio->storage == VMA_IO_PACKED_RAW && + compressed_size != block_bytes) { + pr_err("Packed-raw VMA IO block %d has size %u, expected %llu\n", + i, compressed_size, + (unsigned long long)block_bytes); + return -1; + } + if (rio->storage == VMA_IO_ZERO && compressed_size) { + pr_err("Zero VMA IO block %d has payload size %u\n", i, + compressed_size); + return -1; + } + if (payload_bytes > UINT64_MAX - compressed_size || + output_bytes > UINT64_MAX - block_bytes) { + pr_err("Direct VMA IO size overflows\n"); + return -1; + } + payload_bytes += compressed_size; + output_bytes += block_bytes; + } + + for (i = 0; i < rio->nr_iovs; i++) { + if (iov_bytes > UINT64_MAX - rio->iovs[i].iov_len) { + pr_err("Direct VMA IO destination size overflows\n"); + return -1; + } + iov_bytes += rio->iovs[i].iov_len; + } + + if (payload_bytes != rio->total_compressed_size || + output_bytes != iov_bytes || rio->n_pages <= 0 || + output_bytes != (uint64_t)rio->n_pages * PAGE_SIZE) { + pr_err("Inconsistent direct VMA IO sizes: payload=%llu metadata=%llu output=%llu iov=%llu\n", + (unsigned long long)payload_bytes, + (unsigned long long)rio->total_compressed_size, + (unsigned long long)output_bytes, + (unsigned long long)iov_bytes); + return -1; + } + + return 0; +} + +static int restore_vma_preadv_mixed(struct task_restore_args *args) +{ + struct restore_vma_io *rio = args->vma_ios; + unsigned int i; + int ret = -1; + + /* + * Actual LZ4 blocks are restored into premapped VMAs before PIE. A + * delayed VMA can therefore contain only ordinary uncompressed pages, + * packed raw-fallback blocks, or zero blocks, all of which PIE can + * restore with native syscalls. + */ + for (i = 0; i < args->vma_ios_n; i++) { + if (rio->storage == VMA_IO_UNCOMPRESSED) { + if (args->vma_ios_fd == -1) { + pr_err("No pages image fd for uncompressed VMA IO entry\n"); + goto out; + } + + if (restore_vma_preadv_one(args, rio, true)) + goto out; + goto next; + } + if (rio->storage == VMA_IO_PACKED_RAW) { + if (args->vma_ios_fd == -1) { + pr_err("No pages image fd for packed-raw VMA IO entry\n"); + goto out; + } + if (validate_direct_vma_io(rio) || + restore_vma_preadv_one(args, rio, false)) + goto out; + goto next; + } + if (rio->storage == VMA_IO_ZERO) { + int j; + + if (validate_direct_vma_io(rio)) + goto out; + for (j = 0; j < rio->nr_iovs; j++) + memset(rio->iovs[j].iov_base, 0, + rio->iovs[j].iov_len); + goto next; + } + if (rio->storage != VMA_IO_ENCODED) { + pr_err("Unknown VMA IO storage kind %d\n", rio->storage); + goto out; + } + pr_err("Delayed VMA IO unexpectedly contains an LZ4 block\n"); + goto out; + +next: + rio = (struct restore_vma_io *)((char *)rio + RIO_SIZE(rio->nr_iovs)); + } + + ret = 0; +out: + if (args->vma_ios_fd != -1) { + sys_close(args->vma_ios_fd); + args->vma_ios_fd = -1; + } + + return ret; +} + /* * Call preadv() but limit size of the read. Zero `max_to_read` skips the limit. */ @@ -1794,76 +1937,61 @@ static ssize_t preadv_limited(int fd, struct iovec *iovs, int nr, off_t offs, si return ret; } -/* - * Restore private VMA page contents via synchronous preadv(). - * - * Fallback engine selected when the pages image fd does not support - * O_DIRECT (probe_pages_o_direct() returned 0). Reads each page-io - * iovec batch with sys_preadv() and honors --auto-dedup by punching - * holes in the consumed range. Closes args->vma_ios_fd on exit. - */ -static int restore_vma_preadv(struct task_restore_args *args) +static int restore_vma_preadv_one(struct task_restore_args *args, + struct restore_vma_io *rio, + bool allow_dedup) { - struct restore_vma_io *rio; - unsigned int i; + struct iovec *iovs = rio->iovs; + int nr = rio->nr_iovs; + ssize_t r; - rio = args->vma_ios; - for (i = 0; i < args->vma_ios_n; i++) { - struct iovec *iovs = rio->iovs; - int nr = rio->nr_iovs; - ssize_t r; - - while (nr) { - pr_debug("Preadv %lx:%d... (%d iovs)\n", (unsigned long)iovs->iov_base, (int)iovs->iov_len, nr); - /* - * If we're requested to punch holes in the file after reading we do - * it to save memory. Limit the reads then to an arbitrary block size. - */ - r = preadv_limited(args->vma_ios_fd, iovs, nr, rio->off, - args->auto_dedup ? AUTO_DEDUP_OVERHEAD_BYTES : 0); - if (r < 0) { - pr_err("Can't read pages data (%d)\n", (int)r); - return -1; - } - - if (r == 0) { - pr_err("Unexpected EOF reading pages data at offset %ld (%d iovs remaining)\n", - (long)rio->off, nr); - return -1; - } - - pr_debug("`- returned %ld\n", (long)r); - /* If the file is open for writing, then it means we should punch holes - * in it. */ - if (args->auto_dedup) { - int fr = sys_fallocate(args->vma_ios_fd, FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE, - rio->off, r); - if (fr < 0) { - pr_debug("Failed to punch holes with fallocate: %d\n", fr); - } - } - rio->off += r; - /* Advance the iovecs */ - do { - if (iovs->iov_len <= r) { - pr_debug(" `- skip pagemap\n"); - r -= iovs->iov_len; - iovs++; - nr--; - continue; - } - - iovs->iov_base += r; - iovs->iov_len -= r; - break; - } while (nr > 0); + while (nr) { + pr_debug("Preadv %lx:%d... (%d iovs)\n", (unsigned long)iovs->iov_base, (int)iovs->iov_len, nr); + /* + * If we're requested to punch holes in the file after reading we do + * it to save memory. Limit the reads then to an arbitrary block size. + */ + r = preadv_limited(args->vma_ios_fd, iovs, nr, rio->off, + args->auto_dedup && allow_dedup ? + AUTO_DEDUP_OVERHEAD_BYTES : 0); + if (r < 0) { + pr_err("Can't read pages data (%d)\n", (int)r); + return -1; } - rio = (struct restore_vma_io *)((char *)rio + RIO_SIZE(rio->nr_iovs)); - } + if (r == 0) { + pr_err("Unexpected EOF reading pages data at offset %ld (%d iovs remaining)\n", + (long)rio->off, nr); + return -1; + } - if (args->vma_ios_fd != -1) - sys_close(args->vma_ios_fd); + pr_debug("`- returned %ld\n", (long)r); + /* + * If the file is open for writing, then it means we should + * punch holes in it. + */ + if (args->auto_dedup && allow_dedup) { + int fr = sys_fallocate(args->vma_ios_fd, FALLOC_FL_KEEP_SIZE | FALLOC_FL_PUNCH_HOLE, + rio->off, r); + if (fr < 0) + pr_debug("Failed to punch holes with fallocate: %d\n", fr); + } + rio->off += r; + /* Advance the iovecs */ + do { + if (iovs->iov_len <= r) { + pr_debug(" `- skip pagemap\n"); + r -= iovs->iov_len; + iovs++; + nr--; + continue; + } + + iovs->iov_base += r; + iovs->iov_len -= r; + break; + } while (nr > 0); + } return 0; } @@ -1871,10 +1999,10 @@ static int restore_vma_preadv(struct task_restore_args *args) /* * Restore private VMA page contents via Linux native AIO. * - * Submits io_submit() requests against the O_DIRECT-prepared pages image - * fd in a bounded AIO_BATCH window. Handles MAX_RW_COUNT (0x7FFFF000) - * short reads by masking to page alignment and resubmitting. Closes - * args->vma_ios_fd on exit. + * Submits aligned uncompressed and packed-raw requests against the + * O_DIRECT-prepared pages image in a bounded AIO_BATCH window. Zero records + * are cleared without I/O. Handles MAX_RW_COUNT (0x7FFFF000) short reads by + * masking to page alignment and resubmitting. Closes args->vma_ios_fd on exit. */ static int restore_vma_aio(struct task_restore_args *args) { @@ -1888,40 +2016,75 @@ static int restore_vma_aio(struct task_restore_args *args) struct restore_vma_io *rio; struct io_event *events; unsigned long alloc_sz; - unsigned int submitted = 0, completed = 0; + const unsigned long event_sz = + AIO_BATCH * sizeof(struct io_event); + const unsigned long request_sz = sizeof(struct iocb) + + sizeof(struct iocb *) + sizeof(struct restore_vma_io *); + unsigned int submitted = 0, completed = 0, read_count = 0; unsigned int i; int ret = -1; + if (__builtin_mul_overflow((unsigned long)n, request_sz, &alloc_sz) || + __builtin_add_overflow(alloc_sz, event_sz, &alloc_sz)) { + pr_err("AIO restore metadata size overflows for %u requests\n", n); + sys_close(fd); + args->vma_ios_fd = -1; + return -1; + } + aio_ret = sys_io_setup(AIO_BATCH, &aio_ctx); if (aio_ret < 0) { pr_err("io_setup(%d) failed: %ld\n", AIO_BATCH, aio_ret); + sys_close(fd); + args->vma_ios_fd = -1; return -1; } - alloc_sz = n * sizeof(struct iocb) + n * sizeof(struct iocb *) + - n * sizeof(struct restore_vma_io *) + - AIO_BATCH * sizeof(struct io_event); iocbs = (void *)sys_mmap(NULL, alloc_sz, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); if (IS_ERR(iocbs)) { pr_err("Can't mmap AIO buffers: %ld\n", PTR_ERR(iocbs)); sys_io_destroy(aio_ctx); + sys_close(fd); + args->vma_ios_fd = -1; return -1; } iocbps = (struct iocb **)((char *)iocbs + n * sizeof(struct iocb)); rio_ptrs = (struct restore_vma_io **)((char *)iocbps + n * sizeof(struct iocb *)); events = (struct io_event *)((char *)rio_ptrs + n * sizeof(struct restore_vma_io *)); - /* Build all iocbs from vma_ios */ + /* Build compact iocbs for entries which actually read from the image. */ rio = args->vma_ios; for (i = 0; i < n; i++) { - struct iocb *cb = &iocbs[i]; + struct iocb *cb; size_t expected = 0; int j; - for (j = 0; j < rio->nr_iovs; j++) - expected += rio->iovs[j].iov_len; + if (rio->storage == VMA_IO_ZERO) { + if (validate_direct_vma_io(rio)) + goto out; + for (j = 0; j < rio->nr_iovs; j++) + memset(rio->iovs[j].iov_base, 0, + rio->iovs[j].iov_len); + goto next; + } + if (rio->storage == VMA_IO_PACKED_RAW) { + if (validate_direct_vma_io(rio)) + goto out; + } else if (rio->storage != VMA_IO_UNCOMPRESSED) { + pr_err("AIO restore received storage kind %d\n", rio->storage); + goto out; + } + for (j = 0; j < rio->nr_iovs; j++) { + if (expected > SIZE_MAX - rio->iovs[j].iov_len) { + pr_err("AIO restore request size overflows\n"); + goto out; + } + expected += rio->iovs[j].iov_len; + } + + cb = &iocbs[read_count]; memset(cb, 0, sizeof(*cb)); cb->aio_fildes = fd; cb->aio_lio_opcode = IOCB_CMD_PREADV; @@ -1930,18 +2093,21 @@ static int restore_vma_aio(struct task_restore_args *args) cb->aio_offset = rio->off; /* io_getevents() returns this as event.data for short-read checks. */ cb->aio_data = expected; - iocbps[i] = cb; - rio_ptrs[i] = rio; + iocbps[read_count] = cb; + rio_ptrs[read_count] = rio; + read_count++; + next: rio = (struct restore_vma_io *)((char *)rio + RIO_SIZE(rio->nr_iovs)); } /* Submit and reap in batches */ - while (submitted < n || completed < n) { - if (submit_aio_batch(aio_ctx, iocbps, n, &submitted, completed) < 0) + while (submitted < read_count || completed < read_count) { + if (submit_aio_batch(aio_ctx, iocbps, read_count, + &submitted, completed) < 0) goto out; - if (completed >= n) + if (completed >= read_count) continue; /* @@ -2257,13 +2423,20 @@ __visible long __export_restore_task(struct task_restore_args *args) * Engine is selected from args->vma_ios_use_direct, populated by * probe_pages_o_direct() at restore-args build time: * true -> restore_vma_aio() (io_submit, O_DIRECT, batched) - * false -> restore_vma_preadv() (sys_preadv, buffered, sequential) - * Both helpers consume args->vma_ios and close args->vma_ios_fd on exit. + * false -> restore_vma_preadv_mixed() (buffered, sequential I/O) + * The selected reader consumes args->vma_ios and closes + * args->vma_ios_fd on exit. Compressed images can use AIO when every + * delayed file-backed range is raw and page-aligned; encoded ranges are + * premapped and decoded before PIE. */ if (args->vma_ios_n > 0 && args->vma_ios_fd != -1) { - int rc = args->vma_ios_use_direct - ? restore_vma_aio(args) - : restore_vma_preadv(args); + int rc; + + if (args->vma_ios_use_direct) + pr_debug("Restoring delayed VMA I/O with native AIO\n"); + rc = args->vma_ios_use_direct ? + restore_vma_aio(args) : + restore_vma_preadv_mixed(args); if (rc < 0) goto core_restore_end; }