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; }