// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright 2023 Red Hat
 */
#include "volume-index.h"

#include <linux/bitops.h>
#include <linux/bits.h>
#include <linux/cache.h>
#include <linux/compiler.h>
#include <linux/log2.h>

#include "errors.h"
#include "logger.h"
#include "memory-alloc.h"
#include "numeric.h"
#include "permassert.h"
#include "thread-utils.h"

#include "config.h"
#include "geometry.h"
#include "hash-utils.h"
#include "indexer.h"

/*
 * The volume index is a combination of two separate subindexes, one containing sparse hook entries
 * (retained for all chapters), and one containing the remaining entries (retained only for the
 * dense chapters). If there are no sparse chapters, only the non-hook sub index is used, and it
 * will contain all records for all chapters.
 *
 * The volume index is also divided into zones, with one thread operating on each zone. Each
 * incoming request is dispatched to the appropriate thread, and then to the appropriate subindex.
 * Each delta list is handled by a single zone. To ensure that the distribution of delta lists to
 * zones doesn't underflow (leaving some zone with no delta lists), the minimum number of delta
 * lists must be the square of the maximum zone count for both subindexes.
 *
 * Each subindex zone is a delta index where the payload is a chapter number. The volume index can
 * compute the delta list number, address, and zone number from the record name in order to
 * dispatch record handling to the correct structures.
 *
 * Most operations that use all the zones take place either before request processing is allowed,
 * or after all requests have been flushed in order to shut down. The only multi-threaded operation
 * supported during normal operation is the uds_lookup_volume_index_name() method, used to determine
 * whether a new chapter should be loaded into the sparse index cache. This operation only uses the
 * sparse hook subindex, and the zone mutexes are used to make this operation safe.
 *
 * There are three ways of expressing chapter numbers in the volume index: virtual, index, and
 * rolling. The interface to the volume index uses virtual chapter numbers, which are 64 bits long.
 * Internally the subindex stores only the minimal number of bits necessary by masking away the
 * high-order bits. When the index needs to deal with ordering of index chapter numbers, as when
 * flushing entries from older chapters, it rolls the index chapter number around so that the
 * smallest one in use is mapped to 0. See convert_index_to_virtual() or flush_invalid_entries()
 * for an example of this technique.
 *
 * For efficiency, when older chapter numbers become invalid, the index does not immediately remove
 * the invalidated entries. Instead it lazily removes them from a given delta list the next time it
 * walks that list during normal operation. Because of this, the index size must be increased
 * somewhat to accommodate all the invalid entries that have not yet been removed. For the standard
 * index sizes, this requires about 4 chapters of old entries per 1024 chapters of valid entries in
 * the index.
 */

struct sub_index_parameters {
	/* The number of bits in address mask */
	u8 address_bits;
	/* The number of bits in chapter number */
	u8 chapter_bits;
	/* The mean delta */
	u32 mean_delta;
	/* The number of delta lists */
	u64 list_count;
	/* The number of chapters used */
	u32 chapter_count;
	/* The number of bits per chapter */
	size_t chapter_size_in_bits;
	/* The number of bytes of delta list memory */
	size_t memory_size;
	/* The number of bytes the index should keep free at all times */
	size_t target_free_bytes;
};

struct split_config {
	/* The hook subindex configuration */
	struct uds_configuration hook_config;
	struct index_geometry hook_geometry;

	/* The non-hook subindex configuration */
	struct uds_configuration non_hook_config;
	struct index_geometry non_hook_geometry;
};

struct chapter_range {
	u32 chapter_start;
	u32 chapter_count;
};

#define MAGIC_SIZE 8

static const char MAGIC_START_5[] = "MI5-0005";

struct sub_index_data {
	char magic[MAGIC_SIZE]; /* MAGIC_START_5 */
	u64 volume_nonce;
	u64 virtual_chapter_low;
	u64 virtual_chapter_high;
	u32 first_list;
	u32 list_count;
};

static const char MAGIC_START_6[] = "MI6-0001";

struct volume_index_data {
	char magic[MAGIC_SIZE]; /* MAGIC_START_6 */
	u32 sparse_sample_rate;
};

static inline u32 extract_address(const struct volume_sub_index *sub_index,
				  const struct uds_record_name *name)
{
	return uds_extract_volume_index_bytes(name) & sub_index->address_mask;
}

static inline u32 extract_dlist_num(const struct volume_sub_index *sub_index,
				    const struct uds_record_name *name)
{
	u64 bits = uds_extract_volume_index_bytes(name);

	return (bits >> sub_index->address_bits) % sub_index->list_count;
}

static inline const struct volume_sub_index_zone *
get_zone_for_record(const struct volume_index_record *record)
{
	return &record->sub_index->zones[record->zone_number];
}

static inline u64 convert_index_to_virtual(const struct volume_index_record *record,
					   u32 index_chapter)
{
	const struct volume_sub_index_zone *volume_index_zone = get_zone_for_record(record);
	u32 rolling_chapter = ((index_chapter - volume_index_zone->virtual_chapter_low) &
			       record->sub_index->chapter_mask);

	return volume_index_zone->virtual_chapter_low + rolling_chapter;
}

static inline u32 convert_virtual_to_index(const struct volume_sub_index *sub_index,
					   u64 virtual_chapter)
{
	return virtual_chapter & sub_index->chapter_mask;
}

static inline bool is_virtual_chapter_indexed(const struct volume_index_record *record,
					      u64 virtual_chapter)
{
	const struct volume_sub_index_zone *volume_index_zone = get_zone_for_record(record);

	return ((virtual_chapter >= volume_index_zone->virtual_chapter_low) &&
		(virtual_chapter <= volume_index_zone->virtual_chapter_high));
}

static inline bool has_sparse(const struct volume_index *volume_index)
{
	return volume_index->sparse_sample_rate > 0;
}

bool uds_is_volume_index_sample(const struct volume_index *volume_index,
				const struct uds_record_name *name)
{
	if (!has_sparse(volume_index))
		return false;

	return (uds_extract_sampling_bytes(name) % volume_index->sparse_sample_rate) == 0;
}

static inline const struct volume_sub_index *
get_volume_sub_index(const struct volume_index *volume_index,
		     const struct uds_record_name *name)
{
	return (uds_is_volume_index_sample(volume_index, name) ?
		&volume_index->vi_hook :
		&volume_index->vi_non_hook);
}

static unsigned int get_volume_sub_index_zone(const struct volume_sub_index *sub_index,
					      const struct uds_record_name *name)
{
	return extract_dlist_num(sub_index, name) / sub_index->delta_index.lists_per_zone;
}

unsigned int uds_get_volume_index_zone(const struct volume_index *volume_index,
				       const struct uds_record_name *name)
{
	return get_volume_sub_index_zone(get_volume_sub_index(volume_index, name), name);
}

#define DELTA_LIST_SIZE 256

static int compute_volume_sub_index_parameters(const struct uds_configuration *config,
					       struct sub_index_parameters *params)
{
	u64 entries_in_volume_index, address_span;
	u32 chapters_in_volume_index, invalid_chapters;
	u32 rounded_chapters;
	u64 delta_list_records;
	u32 address_count;
	u64 index_size_in_bits;
	size_t expected_index_size;
	u64 min_delta_lists = MAX_ZONES * MAX_ZONES;
	struct index_geometry *geometry = config->geometry;
	u64 records_per_chapter = geometry->records_per_chapter;

	params->chapter_count = geometry->chapters_per_volume;
	/*
	 * Make sure that the number of delta list records in the volume index does not change when
	 * the volume is reduced by one chapter. This preserves the mapping from name to volume
	 * index delta list.
	 */
	rounded_chapters = params->chapter_count;
	if (uds_is_reduced_index_geometry(geometry))
		rounded_chapters += 1;
	delta_list_records = records_per_chapter * rounded_chapters;
	address_count = config->volume_index_mean_delta * DELTA_LIST_SIZE;
	params->list_count = max(delta_list_records / DELTA_LIST_SIZE, min_delta_lists);
	params->address_bits = bits_per(address_count - 1);
	params->chapter_bits = bits_per(rounded_chapters - 1);
	if ((u32) params->list_count != params->list_count) {
		return vdo_log_warning_strerror(UDS_INVALID_ARGUMENT,
						"cannot initialize volume index with %llu delta lists",
						(unsigned long long) params->list_count);
	}

	if (params->address_bits > 31) {
		return vdo_log_warning_strerror(UDS_INVALID_ARGUMENT,
						"cannot initialize volume index with %u address bits",
						params->address_bits);
	}

	/*
	 * The probability that a given delta list is not touched during the writing of an entire
	 * chapter is:
	 *
	 * double p_not_touched = pow((double) (params->list_count - 1) / params->list_count,
	 *                            records_per_chapter);
	 *
	 * For the standard index sizes, about 78% of the delta lists are not touched, and
	 * therefore contain old index entries that have not been eliminated by the lazy LRU
	 * processing. Then the number of old index entries that accumulate over the entire index,
	 * in terms of full chapters worth of entries, is:
	 *
	 * double invalid_chapters = p_not_touched / (1.0 - p_not_touched);
	 *
	 * For the standard index sizes, the index needs about 3.5 chapters of space for the old
	 * entries in a 1024 chapter index, so round this up to use 4 chapters per 1024 chapters in
	 * the index.
	 */
	invalid_chapters = max(rounded_chapters / 256, 2U);
	chapters_in_volume_index = rounded_chapters + invalid_chapters;
	entries_in_volume_index = records_per_chapter * chapters_in_volume_index;

	address_span = params->list_count << params->address_bits;
	params->mean_delta = address_span / entries_in_volume_index;

	/*
	 * Compute the expected size of a full index, then set the total memory to be 6% larger
	 * than that expected size. This number should be large enough that there are not many
	 * rebalances when the index is full.
	 */
	params->chapter_size_in_bits = uds_compute_delta_index_size(records_per_chapter,
								    params->mean_delta,
								    params->chapter_bits);
	index_size_in_bits = params->chapter_size_in_bits * chapters_in_volume_index;
	expected_index_size = index_size_in_bits / BITS_PER_BYTE;
	params->memory_size = expected_index_size * 106 / 100;

	params->target_free_bytes = expected_index_size / 20;
	return UDS_SUCCESS;
}

static void uninitialize_volume_sub_index(struct volume_sub_index *sub_index)
{
	vdo_free(vdo_forget(sub_index->flush_chapters));
	vdo_free(vdo_forget(sub_index->zones));
	uds_uninitialize_delta_index(&sub_index->delta_index);
}

void uds_free_volume_index(struct volume_index *volume_index)
{
	if (volume_index == NULL)
		return;

	if (volume_index->zones != NULL) {
		unsigned int zone;

		for (zone = 0; zone < volume_index->zone_count; zone++)
			mutex_destroy(&volume_index->zones[zone].hook_mutex);
		vdo_free(vdo_forget(volume_index->zones));
	}

	uninitialize_volume_sub_index(&volume_index->vi_non_hook);
	uninitialize_volume_sub_index(&volume_index->vi_hook);
	vdo_free(volume_index);
}


static int compute_volume_sub_index_save_bytes(const struct uds_configuration *config,
					       size_t *bytes)
{
	struct sub_index_parameters params = { .address_bits = 0 };
	int result;

	result = compute_volume_sub_index_parameters(config, &params);
	if (result != UDS_SUCCESS)
		return result;

	*bytes = (sizeof(struct sub_index_data) + params.list_count * sizeof(u64) +
		  uds_compute_delta_index_save_bytes(params.list_count,
						     params.memory_size));
	return UDS_SUCCESS;
}

/* This function is only useful if the configuration includes sparse chapters. */
static void split_configuration(const struct uds_configuration *config,
				struct split_config *split)
{
	u64 sample_rate, sample_records;
	u64 dense_chapters, sparse_chapters;

	/* Start with copies of the base configuration. */
	split->hook_config = *config;
	split->hook_geometry = *config->geometry;
	split->hook_config.geometry = &split->hook_geometry;
	split->non_hook_config = *config;
	split->non_hook_geometry = *config->geometry;
	split->non_hook_config.geometry = &split->non_hook_geometry;

	sample_rate = config->sparse_sample_rate;
	sparse_chapters = config->geometry->sparse_chapters_per_volume;
	dense_chapters = config->geometry->chapters_per_volume - sparse_chapters;
	sample_records = config->geometry->records_per_chapter / sample_rate;

	/* Adjust the number of records indexed for each chapter. */
	split->hook_geometry.records_per_chapter = sample_records;
	split->non_hook_geometry.records_per_chapter -= sample_records;

	/* Adjust the number of chapters indexed. */
	split->hook_geometry.sparse_chapters_per_volume = 0;
	split->non_hook_geometry.sparse_chapters_per_volume = 0;
	split->non_hook_geometry.chapters_per_volume = dense_chapters;
}

static int compute_volume_index_save_bytes(const struct uds_configuration *config,
					   size_t *bytes)
{
	size_t hook_bytes, non_hook_bytes;
	struct split_config split;
	int result;

	if (!uds_is_sparse_index_geometry(config->geometry))
		return compute_volume_sub_index_save_bytes(config, bytes);

	split_configuration(config, &split);
	result = compute_volume_sub_index_save_bytes(&split.hook_config, &hook_bytes);
	if (result != UDS_SUCCESS)
		return result;

	result = compute_volume_sub_index_save_bytes(&split.non_hook_config,
						     &non_hook_bytes);
	if (result != UDS_SUCCESS)
		return result;

	*bytes = sizeof(struct volume_index_data) + hook_bytes + non_hook_bytes;
	return UDS_SUCCESS;
}

int uds_compute_volume_index_save_blocks(const struct uds_configuration *config,
					 size_t block_size, u64 *block_count)
{
	size_t bytes;
	int result;

	result = compute_volume_index_save_bytes(config, &bytes);
	if (result != UDS_SUCCESS)
		return result;

	bytes += sizeof(struct delta_list_save_info);
	*block_count = DIV_ROUND_UP(bytes, block_size) + MAX_ZONES;
	return UDS_SUCCESS;
}

/* Flush invalid entries while walking the delta list. */
static inline int flush_invalid_entries(struct volume_index_record *record,
					struct chapter_range *flush_range,
					u32 *next_chapter_to_invalidate)
{
	int result;

	result = uds_next_delta_index_entry(&record->delta_entry);
	if (result != UDS_SUCCESS)
		return result;

	while (!record->delta_entry.at_end) {
		u32 index_chapter = uds_get_delta_entry_value(&record->delta_entry);
		u32 relative_chapter = ((index_chapter - flush_range->chapter_start) &
					record->sub_index->chapter_mask);

		if (likely(relative_chapter >= flush_range->chapter_count)) {
			if (relative_chapter < *next_chapter_to_invalidate)
				*next_chapter_to_invalidate = relative_chapter;
			break;
		}

		result = uds_remove_delta_index_entry(&record->delta_entry);
		if (result != UDS_SUCCESS)
			return result;
	}

	return UDS_SUCCESS;
}

/* Find the matching record, or the list offset where the record would go. */
static int get_volume_index_entry(struct volume_index_record *record, u32 list_number,
				  u32 key, struct chapter_range *flush_range)
{
	struct volume_index_record other_record;
	const struct volume_sub_index *sub_index = record->sub_index;
	u32 next_chapter_to_invalidate = sub_index->chapter_mask;
	int result;

	result = uds_start_delta_index_search(&sub_index->delta_index, list_number, 0,
					      &record->delta_entry);
	if (result != UDS_SUCCESS)
		return result;

	do {
		result = flush_invalid_entries(record, flush_range,
					       &next_chapter_to_invalidate);
		if (result != UDS_SUCCESS)
			return result;
	} while (!record->delta_entry.at_end && (key > record->delta_entry.key));

	result = uds_remember_delta_index_offset(&record->delta_entry);
	if (result != UDS_SUCCESS)
		return result;

	/* Check any collision records for a more precise match. */
	other_record = *record;
	if (!other_record.delta_entry.at_end && (key == other_record.delta_entry.key)) {
		for (;;) {
			u8 collision_name[UDS_RECORD_NAME_SIZE];

			result = flush_invalid_entries(&other_record, flush_range,
						       &next_chapter_to_invalidate);
			if (result != UDS_SUCCESS)
				return result;

			if (other_record.delta_entry.at_end ||
			    !other_record.delta_entry.is_collision)
				break;

			result = uds_get_delta_entry_collision(&other_record.delta_entry,
							       collision_name);
			if (result != UDS_SUCCESS)
				return result;

			if (memcmp(collision_name, record->name, UDS_RECORD_NAME_SIZE) == 0) {
				*record = other_record;
				break;
			}
		}
	}
	while (!other_record.delta_entry.at_end) {
		result = flush_invalid_entries(&other_record, flush_range,
					       &next_chapter_to_invalidate);
		if (result != UDS_SUCCESS)
			return result;
	}
	next_chapter_to_invalidate += flush_range->chapter_start;
	next_chapter_to_invalidate &= sub_index->chapter_mask;
	flush_range->chapter_start = next_chapter_to_invalidate;
	flush_range->chapter_count = 0;
	return UDS_SUCCESS;
}

static int get_volume_sub_index_record(struct volume_sub_index *sub_index,
				       const struct uds_record_name *name,
				       struct volume_index_record *record)
{
	int result;
	const struct volume_sub_index_zone *volume_index_zone;
	u32 address = extract_address(sub_index, name);
	u32 delta_list_number = extract_dlist_num(sub_index, name);
	u64 flush_chapter = sub_index->flush_chapters[delta_list_number];

	record->sub_index = sub_index;
	record->mutex = NULL;
	record->name = name;
	record->zone_number = delta_list_number / sub_index->delta_index.lists_per_zone;
	volume_index_zone = get_zone_for_record(record);

	if (flush_chapter < volume_index_zone->virtual_chapter_low) {
		struct chapter_range range;
		u64 flush_count = volume_index_zone->virtual_chapter_low - flush_chapter;

		range.chapter_start = convert_virtual_to_index(sub_index, flush_chapter);
		range.chapter_count = (flush_count > sub_index->chapter_mask ?
				       sub_index->chapter_mask + 1 :
				       flush_count);
		result = get_volume_index_entry(record, delta_list_number, address,
						&range);
		flush_chapter = convert_index_to_virtual(record, range.chapter_start);
		if (flush_chapter > volume_index_zone->virtual_chapter_high)
			flush_chapter = volume_index_zone->virtual_chapter_high;
		sub_index->flush_chapters[delta_list_number] = flush_chapter;
	} else {
		result = uds_get_delta_index_entry(&sub_index->delta_index,
						   delta_list_number, address,
						   name->name, &record->delta_entry);
	}

	if (result != UDS_SUCCESS)
		return result;

	record->is_found =
		(!record->delta_entry.at_end && (record->delta_entry.key == address));
	if (record->is_found) {
		u32 index_chapter = uds_get_delta_entry_value(&record->delta_entry);

		record->virtual_chapter = convert_index_to_virtual(record, index_chapter);
	}

	record->is_collision = record->delta_entry.is_collision;
	return UDS_SUCCESS;
}

int uds_get_volume_index_record(struct volume_index *volume_index,
				const struct uds_record_name *name,
				struct volume_index_record *record)
{
	int result;

	if (uds_is_volume_index_sample(volume_index, name)) {
		/*
		 * Other threads cannot be allowed to call uds_lookup_volume_index_name() while
		 * this thread is finding the volume index record. Due to the lazy LRU flushing of
		 * the volume index, uds_get_volume_index_record() is not a read-only operation.
		 */
		unsigned int zone =
			get_volume_sub_index_zone(&volume_index->vi_hook, name);
		struct mutex *mutex = &volume_index->zones[zone].hook_mutex;

		mutex_lock(mutex);
		result = get_volume_sub_index_record(&volume_index->vi_hook, name,
						     record);
		mutex_unlock(mutex);
		/* Remember the mutex so that other operations on the index record can use it. */
		record->mutex = mutex;
	} else {
		result = get_volume_sub_index_record(&volume_index->vi_non_hook, name,
						     record);
	}

	return result;
}

int uds_put_volume_index_record(struct volume_index_record *record, u64 virtual_chapter)
{
	int result;
	u32 address;
	const struct volume_sub_index *sub_index = record->sub_index;

	if (!is_virtual_chapter_indexed(record, virtual_chapter)) {
		u64 low = get_zone_for_record(record)->virtual_chapter_low;
		u64 high = get_zone_for_record(record)->virtual_chapter_high;

		return vdo_log_warning_strerror(UDS_INVALID_ARGUMENT,
						"cannot put record into chapter number %llu that is out of the valid range %llu to %llu",
						(unsigned long long) virtual_chapter,
						(unsigned long long) low,
						(unsigned long long) high);
	}
	address = extract_address(sub_index, record->name);
	if (unlikely(record->mutex != NULL))
		mutex_lock(record->mutex);
	result = uds_put_delta_index_entry(&record->delta_entry, address,
					   convert_virtual_to_index(sub_index,
								    virtual_chapter),
					   record->is_found ? record->name->name : NULL);
	if (unlikely(record->mutex != NULL))
		mutex_unlock(record->mutex);
	switch (result) {
	case UDS_SUCCESS:
		record->virtual_chapter = virtual_chapter;
		record->is_collision = record->delta_entry.is_collision;
		record->is_found = true;
		break;
	case UDS_OVERFLOW:
		vdo_log_ratelimit(vdo_log_warning_strerror, UDS_OVERFLOW,
				  "Volume index entry dropped due to overflow condition");
		uds_log_delta_index_entry(&record->delta_entry);
		break;
	default:
		break;
	}

	return result;
}

int uds_remove_volume_index_record(struct volume_index_record *record)
{
	int result;

	if (!record->is_found)
		return vdo_log_warning_strerror(UDS_BAD_STATE,
						"illegal operation on new record");

	/* Mark the record so that it cannot be used again */
	record->is_found = false;
	if (unlikely(record->mutex != NULL))
		mutex_lock(record->mutex);
	result = uds_remove_delta_index_entry(&record->delta_entry);
	if (unlikely(record->mutex != NULL))
		mutex_unlock(record->mutex);
	return result;
}

static void set_volume_sub_index_zone_open_chapter(struct volume_sub_index *sub_index,
						   unsigned int zone_number,
						   u64 virtual_chapter)
{
	u64 used_bits = 0;
	struct volume_sub_index_zone *zone = &sub_index->zones[zone_number];
	struct delta_zone *delta_zone;
	u32 i;

	zone->virtual_chapter_low = (virtual_chapter >= sub_index->chapter_count ?
				     virtual_chapter - sub_index->chapter_count + 1 :
				     0);
	zone->virtual_chapter_high = virtual_chapter;

	/* Check to see if the new zone data is too large. */
	delta_zone = &sub_index->delta_index.delta_zones[zone_number];
	for (i = 1; i <= delta_zone->list_count; i++)
		used_bits += delta_zone->delta_lists[i].size;

	if (used_bits > sub_index->max_zone_bits) {
		/* Expire enough chapters to free the desired space. */
		u64 expire_count =
			1 + (used_bits - sub_index->max_zone_bits) / sub_index->chapter_zone_bits;

		if (expire_count == 1) {
			vdo_log_ratelimit(vdo_log_info,
					  "zone %u:  At chapter %llu, expiring chapter %llu early",
					  zone_number,
					  (unsigned long long) virtual_chapter,
					  (unsigned long long) zone->virtual_chapter_low);
			zone->early_flushes++;
			zone->virtual_chapter_low++;
		} else {
			u64 first_expired = zone->virtual_chapter_low;

			if (first_expired + expire_count < zone->virtual_chapter_high) {
				zone->early_flushes += expire_count;
				zone->virtual_chapter_low += expire_count;
			} else {
				zone->early_flushes +=
					zone->virtual_chapter_high - zone->virtual_chapter_low;
				zone->virtual_chapter_low = zone->virtual_chapter_high;
			}
			vdo_log_ratelimit(vdo_log_info,
					  "zone %u:  At chapter %llu, expiring chapters %llu to %llu early",
					  zone_number,
					  (unsigned long long) virtual_chapter,
					  (unsigned long long) first_expired,
					  (unsigned long long) zone->virtual_chapter_low - 1);
		}
	}
}

void uds_set_volume_index_zone_open_chapter(struct volume_index *volume_index,
					    unsigned int zone_number,
					    u64 virtual_chapter)
{
	struct mutex *mutex = &volume_index->zones[zone_number].hook_mutex;

	set_volume_sub_index_zone_open_chapter(&volume_index->vi_non_hook, zone_number,
					       virtual_chapter);

	/*
	 * Other threads cannot be allowed to call uds_lookup_volume_index_name() while the open
	 * chapter number is changing.
	 */
	if (has_sparse(volume_index)) {
		mutex_lock(mutex);
		set_volume_sub_index_zone_open_chapter(&volume_index->vi_hook,
						       zone_number, virtual_chapter);
		mutex_unlock(mutex);
	}
}

/*
 * Set the newest open chapter number for the index, while also advancing the oldest valid chapter
 * number.
 */
void uds_set_volume_index_open_chapter(struct volume_index *volume_index,
				       u64 virtual_chapter)
{
	unsigned int zone;

	for (zone = 0; zone < volume_index->zone_count; zone++)
		uds_set_volume_index_zone_open_chapter(volume_index, zone, virtual_chapter);
}

int uds_set_volume_index_record_chapter(struct volume_index_record *record,
					u64 virtual_chapter)
{
	const struct volume_sub_index *sub_index = record->sub_index;
	int result;

	if (!record->is_found)
		return vdo_log_warning_strerror(UDS_BAD_STATE,
						"illegal operation on new record");

	if (!is_virtual_chapter_indexed(record, virtual_chapter)) {
		u64 low = get_zone_for_record(record)->virtual_chapter_low;
		u64 high = get_zone_for_record(record)->virtual_chapter_high;

		return vdo_log_warning_strerror(UDS_INVALID_ARGUMENT,
						"cannot set chapter number %llu that is out of the valid range %llu to %llu",
						(unsigned long long) virtual_chapter,
						(unsigned long long) low,
						(unsigned long long) high);
	}

	if (unlikely(record->mutex != NULL))
		mutex_lock(record->mutex);
	result = uds_set_delta_entry_value(&record->delta_entry,
					   convert_virtual_to_index(sub_index,
								    virtual_chapter));
	if (unlikely(record->mutex != NULL))
		mutex_unlock(record->mutex);
	if (result != UDS_SUCCESS)
		return result;

	record->virtual_chapter = virtual_chapter;
	return UDS_SUCCESS;
}

static u64 lookup_volume_sub_index_name(const struct volume_sub_index *sub_index,
					const struct uds_record_name *name)
{
	int result;
	u32 address = extract_address(sub_index, name);
	u32 delta_list_number = extract_dlist_num(sub_index, name);
	unsigned int zone_number = get_volume_sub_index_zone(sub_index, name);
	const struct volume_sub_index_zone *zone = &sub_index->zones[zone_number];
	u64 virtual_chapter;
	u32 index_chapter;
	u32 rolling_chapter;
	struct delta_index_entry delta_entry;

	result = uds_get_delta_index_entry(&sub_index->delta_index, delta_list_number,
					   address, name->name, &delta_entry);
	if (result != UDS_SUCCESS)
		return NO_CHAPTER;

	if (delta_entry.at_end || (delta_entry.key != address))
		return NO_CHAPTER;

	index_chapter = uds_get_delta_entry_value(&delta_entry);
	rolling_chapter = (index_chapter - zone->virtual_chapter_low) & sub_index->chapter_mask;

	virtual_chapter = zone->virtual_chapter_low + rolling_chapter;
	if (virtual_chapter > zone->virtual_chapter_high)
		return NO_CHAPTER;

	return virtual_chapter;
}

/* Do a read-only lookup of the record name for sparse cache management. */
u64 uds_lookup_volume_index_name(const struct volume_index *volume_index,
				 const struct uds_record_name *name)
{
	unsigned int zone_number = uds_get_volume_index_zone(volume_index, name);
	struct mutex *mutex = &volume_index->zones[zone_number].hook_mutex;
	u64 virtual_chapter;

	if (!uds_is_volume_index_sample(volume_index, name))
		return NO_CHAPTER;

	mutex_lock(mutex);
	virtual_chapter = lookup_volume_sub_index_name(&volume_index->vi_hook, name);
	mutex_unlock(mutex);

	return virtual_chapter;
}

static void abort_restoring_volume_sub_index(struct volume_sub_index *sub_index)
{
	uds_reset_delta_index(&sub_index->delta_index);
}

static void abort_restoring_volume_index(struct volume_index *volume_index)
{
	abort_restoring_volume_sub_index(&volume_index->vi_non_hook);
	if (has_sparse(volume_index))
		abort_restoring_volume_sub_index(&volume_index->vi_hook);
}

static int start_restoring_volume_sub_index(struct volume_sub_index *sub_index,
					    struct buffered_reader **readers,
					    unsigned int reader_count)
{
	unsigned int z;
	int result;
	u64 virtual_chapter_low = 0, virtual_chapter_high = 0;
	unsigned int i;

	for (i = 0; i < reader_count; i++) {
		struct sub_index_data header;
		u8 buffer[sizeof(struct sub_index_data)];
		size_t offset = 0;
		u32 j;

		result = uds_read_from_buffered_reader(readers[i], buffer,
						       sizeof(buffer));
		if (result != UDS_SUCCESS) {
			return vdo_log_warning_strerror(result,
							"failed to read volume index header");
		}

		memcpy(&header.magic, buffer, MAGIC_SIZE);
		offset += MAGIC_SIZE;
		decode_u64_le(buffer, &offset, &header.volume_nonce);
		decode_u64_le(buffer, &offset, &header.virtual_chapter_low);
		decode_u64_le(buffer, &offset, &header.virtual_chapter_high);
		decode_u32_le(buffer, &offset, &header.first_list);
		decode_u32_le(buffer, &offset, &header.list_count);

		result = VDO_ASSERT(offset == sizeof(buffer),
				    "%zu bytes decoded of %zu expected", offset,
				    sizeof(buffer));
		if (result != VDO_SUCCESS)
			result = UDS_CORRUPT_DATA;

		if (memcmp(header.magic, MAGIC_START_5, MAGIC_SIZE) != 0) {
			return vdo_log_warning_strerror(UDS_CORRUPT_DATA,
							"volume index file had bad magic number");
		}

		if (sub_index->volume_nonce == 0) {
			sub_index->volume_nonce = header.volume_nonce;
		} else if (header.volume_nonce != sub_index->volume_nonce) {
			return vdo_log_warning_strerror(UDS_CORRUPT_DATA,
							"volume index volume nonce incorrect");
		}

		if (i == 0) {
			virtual_chapter_low = header.virtual_chapter_low;
			virtual_chapter_high = header.virtual_chapter_high;
		} else if (virtual_chapter_high != header.virtual_chapter_high) {
			u64 low = header.virtual_chapter_low;
			u64 high = header.virtual_chapter_high;

			return vdo_log_warning_strerror(UDS_CORRUPT_DATA,
							"Inconsistent volume index zone files: Chapter range is [%llu,%llu], chapter range %d is [%llu,%llu]",
							(unsigned long long) virtual_chapter_low,
							(unsigned long long) virtual_chapter_high,
							i, (unsigned long long) low,
							(unsigned long long) high);
		} else if (virtual_chapter_low < header.virtual_chapter_low) {
			virtual_chapter_low = header.virtual_chapter_low;
		}

		for (j = 0; j < header.list_count; j++) {
			u8 decoded[sizeof(u64)];

			result = uds_read_from_buffered_reader(readers[i], decoded,
							       sizeof(u64));
			if (result != UDS_SUCCESS) {
				return vdo_log_warning_strerror(result,
								"failed to read volume index flush ranges");
			}

			sub_index->flush_chapters[header.first_list + j] =
				get_unaligned_le64(decoded);
		}
	}

	for (z = 0; z < sub_index->zone_count; z++) {
		memset(&sub_index->zones[z], 0, sizeof(struct volume_sub_index_zone));
		sub_index->zones[z].virtual_chapter_low = virtual_chapter_low;
		sub_index->zones[z].virtual_chapter_high = virtual_chapter_high;
	}

	result = uds_start_restoring_delta_index(&sub_index->delta_index, readers,
						 reader_count);
	if (result != UDS_SUCCESS)
		return vdo_log_warning_strerror(result, "restoring delta index failed");

	return UDS_SUCCESS;
}

static int start_restoring_volume_index(struct volume_index *volume_index,
					struct buffered_reader **buffered_readers,
					unsigned int reader_count)
{
	unsigned int i;
	int result;

	if (!has_sparse(volume_index)) {
		return start_restoring_volume_sub_index(&volume_index->vi_non_hook,
							buffered_readers, reader_count);
	}

	for (i = 0; i < reader_count; i++) {
		struct volume_index_data header;
		u8 buffer[sizeof(struct volume_index_data)];
		size_t offset = 0;

		result = uds_read_from_buffered_reader(buffered_readers[i], buffer,
						       sizeof(buffer));
		if (result != UDS_SUCCESS) {
			return vdo_log_warning_strerror(result,
							"failed to read volume index header");
		}

		memcpy(&header.magic, buffer, MAGIC_SIZE);
		offset += MAGIC_SIZE;
		decode_u32_le(buffer, &offset, &header.sparse_sample_rate);

		result = VDO_ASSERT(offset == sizeof(buffer),
				    "%zu bytes decoded of %zu expected", offset,
				    sizeof(buffer));
		if (result != VDO_SUCCESS)
			result = UDS_CORRUPT_DATA;

		if (memcmp(header.magic, MAGIC_START_6, MAGIC_SIZE) != 0)
			return vdo_log_warning_strerror(UDS_CORRUPT_DATA,
							"volume index file had bad magic number");

		if (i == 0) {
			volume_index->sparse_sample_rate = header.sparse_sample_rate;
		} else if (volume_index->sparse_sample_rate != header.sparse_sample_rate) {
			vdo_log_warning_strerror(UDS_CORRUPT_DATA,
						 "Inconsistent sparse sample rate in delta index zone files: %u vs. %u",
						 volume_index->sparse_sample_rate,
						 header.sparse_sample_rate);
			return UDS_CORRUPT_DATA;
		}
	}

	result = start_restoring_volume_sub_index(&volume_index->vi_non_hook,
						  buffered_readers, reader_count);
	if (result != UDS_SUCCESS)
		return result;

	return start_restoring_volume_sub_index(&volume_index->vi_hook, buffered_readers,
						reader_count);
}

static int finish_restoring_volume_sub_index(struct volume_sub_index *sub_index,
					     struct buffered_reader **buffered_readers,
					     unsigned int reader_count)
{
	return uds_finish_restoring_delta_index(&sub_index->delta_index,
						buffered_readers, reader_count);
}

static int finish_restoring_volume_index(struct volume_index *volume_index,
					 struct buffered_reader **buffered_readers,
					 unsigned int reader_count)
{
	int result;

	result = finish_restoring_volume_sub_index(&volume_index->vi_non_hook,
						   buffered_readers, reader_count);
	if ((result == UDS_SUCCESS) && has_sparse(volume_index)) {
		result = finish_restoring_volume_sub_index(&volume_index->vi_hook,
							   buffered_readers,
							   reader_count);
	}

	return result;
}

int uds_load_volume_index(struct volume_index *volume_index,
			  struct buffered_reader **readers, unsigned int reader_count)
{
	int result;

	/* Start by reading the header section of the stream. */
	result = start_restoring_volume_index(volume_index, readers, reader_count);
	if (result != UDS_SUCCESS)
		return result;

	result = finish_restoring_volume_index(volume_index, readers, reader_count);
	if (result != UDS_SUCCESS) {
		abort_restoring_volume_index(volume_index);
		return result;
	}

	/* Check the final guard lists to make sure there is no extra data. */
	result = uds_check_guard_delta_lists(readers, reader_count);
	if (result != UDS_SUCCESS)
		abort_restoring_volume_index(volume_index);

	return result;
}

static int start_saving_volume_sub_index(const struct volume_sub_index *sub_index,
					 unsigned int zone_number,
					 struct buffered_writer *buffered_writer)
{
	int result;
	struct volume_sub_index_zone *volume_index_zone = &sub_index->zones[zone_number];
	u32 first_list = sub_index->delta_index.delta_zones[zone_number].first_list;
	u32 list_count = sub_index->delta_index.delta_zones[zone_number].list_count;
	u8 buffer[sizeof(struct sub_index_data)];
	size_t offset = 0;
	u32 i;

	memcpy(buffer, MAGIC_START_5, MAGIC_SIZE);
	offset += MAGIC_SIZE;
	encode_u64_le(buffer, &offset, sub_index->volume_nonce);
	encode_u64_le(buffer, &offset, volume_index_zone->virtual_chapter_low);
	encode_u64_le(buffer, &offset, volume_index_zone->virtual_chapter_high);
	encode_u32_le(buffer, &offset, first_list);
	encode_u32_le(buffer, &offset, list_count);

	result =  VDO_ASSERT(offset == sizeof(struct sub_index_data),
			     "%zu bytes of config written, of %zu expected", offset,
			     sizeof(struct sub_index_data));
	if (result != VDO_SUCCESS)
		return result;

	result = uds_write_to_buffered_writer(buffered_writer, buffer, offset);
	if (result != UDS_SUCCESS)
		return vdo_log_warning_strerror(result,
						"failed to write volume index header");

	for (i = 0; i < list_count; i++) {
		u8 encoded[sizeof(u64)];

		put_unaligned_le64(sub_index->flush_chapters[first_list + i], &encoded);
		result = uds_write_to_buffered_writer(buffered_writer, encoded,
						      sizeof(u64));
		if (result != UDS_SUCCESS) {
			return vdo_log_warning_strerror(result,
							"failed to write volume index flush ranges");
		}
	}

	return uds_start_saving_delta_index(&sub_index->delta_index, zone_number,
					    buffered_writer);
}

static int start_saving_volume_index(const struct volume_index *volume_index,
				     unsigned int zone_number,
				     struct buffered_writer *writer)
{
	u8 buffer[sizeof(struct volume_index_data)];
	size_t offset = 0;
	int result;

	if (!has_sparse(volume_index)) {
		return start_saving_volume_sub_index(&volume_index->vi_non_hook,
						     zone_number, writer);
	}

	memcpy(buffer, MAGIC_START_6, MAGIC_SIZE);
	offset += MAGIC_SIZE;
	encode_u32_le(buffer, &offset, volume_index->sparse_sample_rate);
	result = VDO_ASSERT(offset == sizeof(struct volume_index_data),
			    "%zu bytes of header written, of %zu expected", offset,
			    sizeof(struct volume_index_data));
	if (result != VDO_SUCCESS)
		return result;

	result = uds_write_to_buffered_writer(writer, buffer, offset);
	if (result != UDS_SUCCESS) {
		vdo_log_warning_strerror(result, "failed to write volume index header");
		return result;
	}

	result = start_saving_volume_sub_index(&volume_index->vi_non_hook, zone_number,
					       writer);
	if (result != UDS_SUCCESS)
		return result;

	return start_saving_volume_sub_index(&volume_index->vi_hook, zone_number,
					     writer);
}

static int finish_saving_volume_sub_index(const struct volume_sub_index *sub_index,
					  unsigned int zone_number)
{
	return uds_finish_saving_delta_index(&sub_index->delta_index, zone_number);
}

static int finish_saving_volume_index(const struct volume_index *volume_index,
				      unsigned int zone_number)
{
	int result;

	result = finish_saving_volume_sub_index(&volume_index->vi_non_hook, zone_number);
	if ((result == UDS_SUCCESS) && has_sparse(volume_index))
		result = finish_saving_volume_sub_index(&volume_index->vi_hook, zone_number);
	return result;
}

int uds_save_volume_index(struct volume_index *volume_index,
			  struct buffered_writer **writers, unsigned int writer_count)
{
	int result = UDS_SUCCESS;
	unsigned int zone;

	for (zone = 0; zone < writer_count; zone++) {
		result = start_saving_volume_index(volume_index, zone, writers[zone]);
		if (result != UDS_SUCCESS)
			break;

		result = finish_saving_volume_index(volume_index, zone);
		if (result != UDS_SUCCESS)
			break;

		result = uds_write_guard_delta_list(writers[zone]);
		if (result != UDS_SUCCESS)
			break;

		result = uds_flush_buffered_writer(writers[zone]);
		if (result != UDS_SUCCESS)
			break;
	}

	return result;
}

static void get_volume_sub_index_stats(const struct volume_sub_index *sub_index,
				       struct volume_index_stats *stats)
{
	struct delta_index_stats dis;
	unsigned int z;

	uds_get_delta_index_stats(&sub_index->delta_index, &dis);
	stats->rebalance_time = dis.rebalance_time;
	stats->rebalance_count = dis.rebalance_count;
	stats->record_count = dis.record_count;
	stats->collision_count = dis.collision_count;
	stats->discard_count = dis.discard_count;
	stats->overflow_count = dis.overflow_count;
	stats->delta_lists = dis.list_count;
	stats->early_flushes = 0;
	for (z = 0; z < sub_index->zone_count; z++)
		stats->early_flushes += sub_index->zones[z].early_flushes;
}

void uds_get_volume_index_stats(const struct volume_index *volume_index,
				struct volume_index_stats *stats)
{
	struct volume_index_stats sparse_stats;

	get_volume_sub_index_stats(&volume_index->vi_non_hook, stats);
	if (!has_sparse(volume_index))
		return;

	get_volume_sub_index_stats(&volume_index->vi_hook, &sparse_stats);
	stats->rebalance_time += sparse_stats.rebalance_time;
	stats->rebalance_count += sparse_stats.rebalance_count;
	stats->record_count += sparse_stats.record_count;
	stats->collision_count += sparse_stats.collision_count;
	stats->discard_count += sparse_stats.discard_count;
	stats->overflow_count += sparse_stats.overflow_count;
	stats->delta_lists += sparse_stats.delta_lists;
	stats->early_flushes += sparse_stats.early_flushes;
}

static int initialize_volume_sub_index(const struct uds_configuration *config,
				       u64 volume_nonce, u8 tag,
				       struct volume_sub_index *sub_index)
{
	struct sub_index_parameters params = { .address_bits = 0 };
	unsigned int zone_count = config->zone_count;
	u64 available_bytes = 0;
	unsigned int z;
	int result;

	result = compute_volume_sub_index_parameters(config, &params);
	if (result != UDS_SUCCESS)
		return result;

	sub_index->address_bits = params.address_bits;
	sub_index->address_mask = (1u << params.address_bits) - 1;
	sub_index->chapter_bits = params.chapter_bits;
	sub_index->chapter_mask = (1u << params.chapter_bits) - 1;
	sub_index->chapter_count = params.chapter_count;
	sub_index->list_count = params.list_count;
	sub_index->zone_count = zone_count;
	sub_index->chapter_zone_bits = params.chapter_size_in_bits / zone_count;
	sub_index->volume_nonce = volume_nonce;

	result = uds_initialize_delta_index(&sub_index->delta_index, zone_count,
					    params.list_count, params.mean_delta,
					    params.chapter_bits, params.memory_size,
					    tag);
	if (result != UDS_SUCCESS)
		return result;

	for (z = 0; z < sub_index->delta_index.zone_count; z++)
		available_bytes += sub_index->delta_index.delta_zones[z].size;
	available_bytes -= params.target_free_bytes;
	sub_index->max_zone_bits = (available_bytes * BITS_PER_BYTE) / zone_count;
	sub_index->memory_size = (sub_index->delta_index.memory_size +
				  sizeof(struct volume_sub_index) +
				  (params.list_count * sizeof(u64)) +
				  (zone_count * sizeof(struct volume_sub_index_zone)));

	/* The following arrays are initialized to all zeros. */
	result = vdo_allocate(params.list_count, u64, "first chapter to flush",
			      &sub_index->flush_chapters);
	if (result != VDO_SUCCESS)
		return result;

	return vdo_allocate(zone_count, struct volume_sub_index_zone,
			    "volume index zones", &sub_index->zones);
}

int uds_make_volume_index(const struct uds_configuration *config, u64 volume_nonce,
			  struct volume_index **volume_index_ptr)
{
	struct split_config split;
	unsigned int zone;
	struct volume_index *volume_index;
	int result;

	result = vdo_allocate(1, struct volume_index, "volume index", &volume_index);
	if (result != VDO_SUCCESS)
		return result;

	volume_index->zone_count = config->zone_count;

	if (!uds_is_sparse_index_geometry(config->geometry)) {
		result = initialize_volume_sub_index(config, volume_nonce, 'm',
						     &volume_index->vi_non_hook);
		if (result != UDS_SUCCESS) {
			uds_free_volume_index(volume_index);
			return result;
		}

		volume_index->memory_size = volume_index->vi_non_hook.memory_size;
		*volume_index_ptr = volume_index;
		return UDS_SUCCESS;
	}

	volume_index->sparse_sample_rate = config->sparse_sample_rate;

	result = vdo_allocate(config->zone_count, struct volume_index_zone,
			      "volume index zones", &volume_index->zones);
	if (result != VDO_SUCCESS) {
		uds_free_volume_index(volume_index);
		return result;
	}

	for (zone = 0; zone < config->zone_count; zone++)
		mutex_init(&volume_index->zones[zone].hook_mutex);

	split_configuration(config, &split);
	result = initialize_volume_sub_index(&split.non_hook_config, volume_nonce, 'd',
					     &volume_index->vi_non_hook);
	if (result != UDS_SUCCESS) {
		uds_free_volume_index(volume_index);
		return vdo_log_error_strerror(result,
					      "Error creating non hook volume index");
	}

	result = initialize_volume_sub_index(&split.hook_config, volume_nonce, 's',
					     &volume_index->vi_hook);
	if (result != UDS_SUCCESS) {
		uds_free_volume_index(volume_index);
		return vdo_log_error_strerror(result,
					      "Error creating hook volume index");
	}

	volume_index->memory_size =
		volume_index->vi_non_hook.memory_size + volume_index->vi_hook.memory_size;
	*volume_index_ptr = volume_index;
	return UDS_SUCCESS;
}
