httpfs2_extension.cpp

#define DUCKDB_EXTENSION_MAIN

#include "duckdb.hpp"

#include "duckdb/main/setting_info.hpp"
#include "duckdb/planner/expression/bound_function_expression.hpp"

#include "httpfs2_bench.hpp"
#include "httpfs2_cost_model.hpp"
#include "httpfs2_extension.hpp"
#include "httpfs2_fake_s3_filesystem.hpp"
#include "httpfs2_filesystem_wrapper.hpp"
#include "httpfs2_state.hpp"

namespace duckdb {

static uint64_t CeilDivU64(uint64_t num, uint64_t denom) {
	return denom == 0 ? 0 : (num + denom - 1) / denom;
}

static uint64_t Httpfs2ChooseFanout(uint64_t bytes, uint64_t max_fanout, uint64_t overhead_ms) {
	if (max_fanout == 0) {
		throw InvalidInputException("httpfs2_choose_fanout: max_fanout must be > 0");
	}
	if (bytes == 0) {
		return 1;
	}
	uint64_t best_k = 1;
	auto best_cost = static_cast<double>(Httpfs2EstimateS3LatencyMs(bytes));
	for (uint64_t k = 2; k <= max_fanout; k++) {
		const auto per_req = CeilDivU64(bytes, k);
		const auto cost = static_cast<double>(Httpfs2EstimateS3LatencyMs(per_req)) + static_cast<double>((k - 1) * overhead_ms);
		if (cost < best_cost) {
			best_cost = cost;
			best_k = k;
		}
	}
	return best_k;
}

static void Httpfs2EstimateS3LatencyMsScalar(DataChunk &args, ExpressionState &, Vector &result) {
	D_ASSERT(args.ColumnCount() == 1);
	const auto count = args.size();

	UnifiedVectorFormat input_data;
	args.data[0].ToUnifiedFormat(count, input_data);
	auto input_ptr = UnifiedVectorFormat::GetData<uint64_t>(input_data);

	result.SetVectorType(VectorType::FLAT_VECTOR);
	auto out_ptr = FlatVector::GetData<uint64_t>(result);
	auto &validity = FlatVector::Validity(result);

	for (idx_t i = 0; i < count; i++) {
		const auto idx = input_data.sel->get_index(i);
		if (!input_data.validity.RowIsValid(idx)) {
			validity.SetInvalid(i);
			continue;
		}
		out_ptr[i] = Httpfs2EstimateS3LatencyMs(input_ptr[idx]);
	}
}

static void Httpfs2ChooseFanoutScalar(DataChunk &args, ExpressionState &, Vector &result) {
	D_ASSERT(args.ColumnCount() == 3);
	const auto count = args.size();

	UnifiedVectorFormat bytes_data;
	UnifiedVectorFormat max_fanout_data;
	UnifiedVectorFormat overhead_data;
	args.data[0].ToUnifiedFormat(count, bytes_data);
	args.data[1].ToUnifiedFormat(count, max_fanout_data);
	args.data[2].ToUnifiedFormat(count, overhead_data);

	auto bytes_ptr = UnifiedVectorFormat::GetData<uint64_t>(bytes_data);
	auto max_fanout_ptr = UnifiedVectorFormat::GetData<uint64_t>(max_fanout_data);
	auto overhead_ptr = UnifiedVectorFormat::GetData<uint64_t>(overhead_data);

	result.SetVectorType(VectorType::FLAT_VECTOR);
	auto out_ptr = FlatVector::GetData<uint64_t>(result);
	auto &validity = FlatVector::Validity(result);

	for (idx_t i = 0; i < count; i++) {
		const auto idx0 = bytes_data.sel->get_index(i);
		const auto idx1 = max_fanout_data.sel->get_index(i);
		const auto idx2 = overhead_data.sel->get_index(i);
		if (!bytes_data.validity.RowIsValid(idx0) || !max_fanout_data.validity.RowIsValid(idx1) ||
		    !overhead_data.validity.RowIsValid(idx2)) {
			validity.SetInvalid(i);
			continue;
		}
		out_ptr[i] = Httpfs2ChooseFanout(bytes_ptr[idx0], max_fanout_ptr[idx1], overhead_ptr[idx2]);
	}
}

static void Httpfs2WrapFileSystem(DatabaseInstance &instance) {
	auto &config = DBConfig::GetConfig(instance);
	if (dynamic_cast<Httpfs2FileSystemWrapper *>(config.file_system.get())) {
		return;
	}
	auto state = Httpfs2GetOrCreateState(instance);
	config.file_system = make_uniq<Httpfs2FileSystemWrapper>(std::move(config.file_system), std::move(state));
}

static shared_ptr<Httpfs2State> GetStateFromContext(ClientContext &context) {
	if (!context.db) {
		throw InternalException("httpfs2: missing DatabaseInstance");
	}
	return Httpfs2GetOrCreateState(*context.db);
}

static void UpdateEnabled(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.enabled.store(parameter.GetValue<bool>());
}

static void UpdateIOThreads(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	const auto v = parameter.GetValue<uint64_t>();
	if (v == 0) {
		throw InvalidInputException("httpfs2_parallel_read_io_threads must be > 0");
	}
	state->SetIOThreads(v);
}

static void UpdateChunkSize(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	const auto v = parameter.GetValue<uint64_t>();
	if (v == 0) {
		throw InvalidInputException("httpfs2_parallel_read_chunk_size_bytes must be > 0");
	}
	state->config.chunk_size_bytes.store(v);
}

static void UpdateMaxFanout(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	const auto v = parameter.GetValue<uint64_t>();
	if (v == 0) {
		throw InvalidInputException("httpfs2_parallel_read_max_fanout must be > 0");
	}
	state->config.max_fanout.store(v);
}

static void UpdateMinBytes(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.min_bytes.store(parameter.GetValue<uint64_t>());
}

static void UpdateDedicatedHandles(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.dedicated_handles.store(parameter.GetValue<bool>());
}

static void UpdateRecommendedPresetEnabled(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.recommended_preset_enabled.store(parameter.GetValue<bool>());
}

static void UpdateAutoFanoutEnabled(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.auto_fanout_enabled.store(parameter.GetValue<bool>());
}

static void UpdateAutoFanoutOverheadMs(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.auto_fanout_overhead_ms.store(parameter.GetValue<uint64_t>());
}

static void UpdateAutoFanoutOverheadAdaptiveEnabled(ClientContext &context, SetScope scope, Value &parameter) {
	auto state = GetStateFromContext(context);
	state->config.auto_fanout_overhead_adaptive_enabled.store(parameter.GetValue<bool>());
}

static void UpdateFakeS3LatencyEnabled(ClientContext &, SetScope, Value &parameter) {
	Httpfs2FakeS3FileSystem::SetLatencyEnabled(parameter.GetValue<bool>());
}

static void UpdateFakeS3LatencyScalePercent(ClientContext &, SetScope, Value &parameter) {
	Httpfs2FakeS3FileSystem::SetLatencyScalePercent(parameter.GetValue<uint64_t>());
}

static void UpdateFakeS3LatencyExtraMs(ClientContext &, SetScope, Value &parameter) {
	Httpfs2FakeS3FileSystem::SetLatencyExtraMs(parameter.GetValue<uint64_t>());
}

struct Httpfs2StateBindData : public FunctionData {
	shared_ptr<Httpfs2State> state;

	unique_ptr<FunctionData> Copy() const override {
		auto copy = make_uniq<Httpfs2StateBindData>();
		copy->state = state;
		return std::move(copy);
	}

	bool Equals(const FunctionData &other_p) const override {
		auto &other = other_p.Cast<Httpfs2StateBindData>();
		return state.get() == other.state.get();
	}

	static Httpfs2StateBindData &GetFrom(ExpressionState &state) {
		auto &func_expr = state.expr.Cast<BoundFunctionExpression>();
		return func_expr.bind_info->Cast<Httpfs2StateBindData>();
	}
};

static unique_ptr<FunctionData> Httpfs2StateBind(ClientContext &context, ScalarFunction &,
                                                 vector<unique_ptr<Expression>> &) {
	auto bind = make_uniq<Httpfs2StateBindData>();
	bind->state = GetStateFromContext(context);
	return std::move(bind);
}

static void Httpfs2AutoFanoutObservedOverheadMsScalar(DataChunk &, ExpressionState &state, Vector &result) {
	auto &bind = Httpfs2StateBindData::GetFrom(state);
	const uint64_t ms = bind.state ? bind.state->observed_overhead_ms.load() : 0;
	result.SetVectorType(VectorType::CONSTANT_VECTOR);
	ConstantVector::GetData<uint64_t>(result)[0] = ms;
}

void Httpfs2Extension::Load(ExtensionLoader &loader) {
	auto &instance = loader.GetDatabaseInstance();
	Httpfs2GetOrCreateState(instance);
	Httpfs2WrapFileSystem(instance);

	auto &config = DBConfig::GetConfig(instance);
	config.file_system->RegisterSubSystem(make_uniq<Httpfs2FakeS3FileSystem>());

	loader.RegisterFunction(Httpfs2BenchReadFunction());
	loader.RegisterFunction(Httpfs2BenchReadManyFunction());
	loader.RegisterFunction(Httpfs2BenchReadManyPoolFunction());

	loader.RegisterFunction(ScalarFunction("httpfs2_estimate_s3_latency_ms", {LogicalType::UBIGINT}, LogicalType::UBIGINT,
	                                       Httpfs2EstimateS3LatencyMsScalar));
	loader.RegisterFunction(ScalarFunction("httpfs2_choose_fanout",
	                                       {LogicalType::UBIGINT, LogicalType::UBIGINT, LogicalType::UBIGINT},
	                                       LogicalType::UBIGINT, Httpfs2ChooseFanoutScalar));

	ScalarFunction observed_overhead_fun({}, LogicalType::UBIGINT, Httpfs2AutoFanoutObservedOverheadMsScalar,
	                                     Httpfs2StateBind, nullptr, nullptr, nullptr,
	                                     LogicalType(LogicalTypeId::INVALID), FunctionStability::VOLATILE);
	observed_overhead_fun.name = "httpfs2_auto_fanout_observed_overhead_ms";
	loader.RegisterFunction(observed_overhead_fun);

	config.AddExtensionOption("httpfs2_parallel_read_enabled", "Enable httpfs2 parallel reads for remote files.",
	                          LogicalTypeId::BOOLEAN, Value::BOOLEAN(false), UpdateEnabled);

	config.AddExtensionOption("httpfs2_parallel_read_io_threads",
	                          "Number of httpfs2 IO worker threads (separate from DuckDB threads).",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(16), UpdateIOThreads);

	config.AddExtensionOption("httpfs2_parallel_read_chunk_size_bytes",
	                          "Target chunk size used when splitting a single Read into parallel sub-reads.",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(16ULL * 1024ULL * 1024ULL), UpdateChunkSize);

	config.AddExtensionOption("httpfs2_parallel_read_max_fanout",
	                          "Max number of parallel sub-reads issued for a single Read call.",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(4), UpdateMaxFanout);

	config.AddExtensionOption("httpfs2_parallel_read_min_bytes",
	                          "Minimum Read size (bytes) before httpfs2 starts splitting into parallel sub-reads.",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(2ULL * 1024ULL * 1024ULL), UpdateMinBytes);

	config.AddExtensionOption("httpfs2_parallel_read_dedicated_handles",
	                          "Use multiple underlying file handles per remote file to reduce internal httpfs contention.",
	                          LogicalTypeId::BOOLEAN, Value::BOOLEAN(false), UpdateDedicatedHandles);

	config.AddExtensionOption("httpfs2_parallel_read_recommended_preset_enabled",
	                          "If enabled, uses the recommended preset: fanout=2, chunk=16MiB, min_bytes=2MiB, no auto-fanout, no dedicated handles.",
	                          LogicalTypeId::BOOLEAN, Value::BOOLEAN(false), UpdateRecommendedPresetEnabled);

	config.AddExtensionOption("httpfs2_parallel_read_auto_fanout_enabled",
	                          "If enabled, httpfs2 chooses a smaller fanout per Read using a simple S3 cost model.",
	                          LogicalTypeId::BOOLEAN, Value::BOOLEAN(false), UpdateAutoFanoutEnabled);
	config.AddExtensionOption("httpfs2_parallel_read_auto_fanout_overhead_ms",
	                          "Penalty (ms) per extra subrequest used by auto fanout selection.",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(500), UpdateAutoFanoutOverheadMs);
	config.AddExtensionOption("httpfs2_parallel_read_auto_fanout_overhead_adaptive_enabled",
	                          "If enabled, httpfs2 learns extra per-request overhead from observed timings and uses it in auto-fanout selection.",
	                          LogicalTypeId::BOOLEAN, Value::BOOLEAN(false), UpdateAutoFanoutOverheadAdaptiveEnabled);

	config.AddExtensionOption("httpfs2_fake_s3_latency_enabled",
	                          "If enabled, fake_s3:// reads sleep to simulate S3-style latency and throughput.",
	                          LogicalTypeId::BOOLEAN, Value::BOOLEAN(true), UpdateFakeS3LatencyEnabled);
	config.AddExtensionOption("httpfs2_fake_s3_latency_scale_percent",
	                          "Scale factor (percent) applied to fake_s3:// simulated latency.",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(100), UpdateFakeS3LatencyScalePercent);
	config.AddExtensionOption("httpfs2_fake_s3_latency_extra_ms",
	                          "Extra constant latency (ms) added to fake_s3:// simulated reads.",
	                          LogicalTypeId::UBIGINT, Value::UBIGINT(0), UpdateFakeS3LatencyExtraMs);
}

} // namespace duckdb

extern "C" {

DUCKDB_CPP_EXTENSION_ENTRY(httpfs2, loader) {
	duckdb::Httpfs2Extension().Load(loader);
}

}