should be it

external/duckdb/test/sql/parallelism/interquery/test_concurrentupdate.cpp

@@ -0,0 +1,213 @@
+#include "catch.hpp"
+#include "duckdb/common/value_operations/value_operations.hpp"
+#include "test_helpers.hpp"
+
+#include <atomic>
+#include <random>
+#include <thread>
+
+using namespace duckdb;
+using namespace std;
+
+namespace test_concurrent_update {
+
+static constexpr int CONCURRENT_UPDATE_TRANSACTION_UPDATE_COUNT = 1000;
+static constexpr int CONCURRENT_UPDATE_TOTAL_ACCOUNTS = 10;
+static constexpr int CONCURRENT_UPDATE_MONEY_PER_ACCOUNT = 10;
+
+TEST_CASE("Single thread update", "[interquery][.]") {
+	duckdb::unique_ptr<MaterializedQueryResult> result;
+	DuckDB db(nullptr);
+	Connection con(db);
+
+	// enable detailed profiling
+	con.Query("PRAGMA enable_profiling");
+	auto detailed_profiling_output = TestCreatePath("detailed_profiling_output");
+	con.Query("PRAGMA profiling_output='" + detailed_profiling_output + "'");
+	con.Query("PRAGMA profiling_mode = detailed");
+
+	// initialize the database
+	con.Query("CREATE TABLE integers(i INTEGER);");
+	int sum = 0;
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TOTAL_ACCOUNTS; i++) {
+		for (size_t j = 0; j < 10; j++) {
+			con.Query("INSERT INTO integers VALUES (" + to_string(j + 1) + ");");
+			sum += j + 1;
+		}
+	}
+
+	// check the sum
+	result = con.Query("SELECT SUM(i) FROM integers");
+	REQUIRE(CHECK_COLUMN(result, 0, {sum}));
+
+	// simple update, we should update INSERT_ELEMENTS elements
+	result = con.Query("UPDATE integers SET i=4 WHERE i=2");
+	REQUIRE(CHECK_COLUMN(result, 0, {CONCURRENT_UPDATE_TOTAL_ACCOUNTS}));
+
+	// check updated sum
+	result = con.Query("SELECT SUM(i) FROM integers");
+	REQUIRE(CHECK_COLUMN(result, 0, {sum + 2 * CONCURRENT_UPDATE_TOTAL_ACCOUNTS}));
+}
+
+atomic<bool> finished_updating;
+static void read_total_balance(DuckDB *db, bool *read_correct) {
+	*read_correct = true;
+	Connection con(*db);
+	while (!finished_updating) {
+		// the total balance should remain constant regardless of updates
+		auto result = con.Query("SELECT SUM(money) FROM accounts");
+		if (!CHECK_COLUMN(result, 0, {CONCURRENT_UPDATE_TOTAL_ACCOUNTS * CONCURRENT_UPDATE_MONEY_PER_ACCOUNT})) {
+			*read_correct = false;
+		}
+	}
+}
+
+TEST_CASE("Concurrent update", "[interquery][.]") {
+	duckdb::unique_ptr<MaterializedQueryResult> result;
+	DuckDB db(nullptr);
+	Connection con(db);
+
+	// enable detailed profiling
+	con.Query("PRAGMA enable_profiling");
+	auto detailed_profiling_output = TestCreatePath("detailed_profiling_output");
+	con.Query("PRAGMA profiling_output='" + detailed_profiling_output + "'");
+	con.Query("PRAGMA profiling_mode = detailed");
+
+	// fixed seed random numbers
+	mt19937 generator;
+	generator.seed(42);
+	uniform_int_distribution<int> account_distribution(0, CONCURRENT_UPDATE_TOTAL_ACCOUNTS - 1);
+	auto random_account = bind(account_distribution, generator);
+
+	uniform_int_distribution<int> amount_distribution(0, CONCURRENT_UPDATE_MONEY_PER_ACCOUNT);
+	auto random_amount = bind(amount_distribution, generator);
+
+	finished_updating = false;
+	// initialize the database
+	con.Query("CREATE TABLE accounts(id INTEGER, money INTEGER)");
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TOTAL_ACCOUNTS; i++) {
+		con.Query("INSERT INTO accounts VALUES (" + to_string(i) + ", " +
+		          to_string(CONCURRENT_UPDATE_MONEY_PER_ACCOUNT) + ");");
+	}
+
+	bool read_correct;
+	// launch separate thread for reading aggregate
+	thread read_thread(read_total_balance, &db, &read_correct);
+
+	// start vigorously updating balances in this thread
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TRANSACTION_UPDATE_COUNT; i++) {
+		int from = random_account();
+		int to = random_account();
+		while (to == from) {
+			to = random_account();
+		}
+		int amount = random_amount();
+
+		REQUIRE_NO_FAIL(con.Query("BEGIN TRANSACTION"));
+		result = con.Query("SELECT money FROM accounts WHERE id=" + to_string(from));
+		Value money_from = result->GetValue(0, 0);
+		result = con.Query("SELECT money FROM accounts WHERE id=" + to_string(to));
+		Value money_to = result->GetValue(0, 0);
+
+		REQUIRE_NO_FAIL(
+		    con.Query("UPDATE accounts SET money = money - " + to_string(amount) + " WHERE id = " + to_string(from)));
+		REQUIRE_NO_FAIL(
+		    con.Query("UPDATE accounts SET money = money + " + to_string(amount) + " WHERE id = " + to_string(to)));
+
+		result = con.Query("SELECT money FROM accounts WHERE id=" + to_string(from));
+		Value new_money_from = result->GetValue(0, 0);
+		result = con.Query("SELECT money FROM accounts WHERE id=" + to_string(to));
+		Value new_money_to = result->GetValue(0, 0);
+
+		Value expected_money_from, expected_money_to;
+
+		expected_money_from = Value::INTEGER(IntegerValue::Get(money_from) - amount);
+		expected_money_to = Value::INTEGER(IntegerValue::Get(money_to) + amount);
+
+		REQUIRE(new_money_from == expected_money_from);
+		REQUIRE(new_money_to == expected_money_to);
+
+		REQUIRE_NO_FAIL(con.Query("COMMIT"));
+	}
+	finished_updating = true;
+	read_thread.join();
+	REQUIRE(read_correct);
+}
+
+static std::atomic<size_t> finished_threads;
+
+static void write_random_numbers_to_account(DuckDB *db, bool *correct, size_t nr) {
+	correct[nr] = true;
+	Connection con(*db);
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TRANSACTION_UPDATE_COUNT; i++) {
+		// just make some changes to the total
+		// the total amount of money after the commit is the same
+		if (con.Query("BEGIN TRANSACTION")->HasError()) {
+			correct[nr] = false;
+		}
+		if (con.Query("UPDATE accounts SET money = money + " + to_string(i * 2) + " WHERE id = " + to_string(nr))
+		        ->HasError()) {
+			correct[nr] = false;
+		}
+		if (con.Query("UPDATE accounts SET money = money - " + to_string(i) + " WHERE id = " + to_string(nr))
+		        ->HasError()) {
+			correct[nr] = false;
+		}
+		if (con.Query("UPDATE accounts SET money = money - " + to_string(i * 2) + " WHERE id = " + to_string(nr))
+		        ->HasError()) {
+			correct[nr] = false;
+		}
+		if (con.Query("UPDATE accounts SET money = money + " + to_string(i) + " WHERE id = " + to_string(nr))
+		        ->HasError()) {
+			correct[nr] = false;
+		}
+		// we test both commit and rollback
+		// the result of both should be the same since the updates have a
+		// net-zero effect
+		if (con.Query(nr % 2 == 0 ? "COMMIT" : "ROLLBACK")->HasError()) {
+			correct[nr] = false;
+		}
+	}
+	finished_threads++;
+	if (finished_threads == CONCURRENT_UPDATE_TOTAL_ACCOUNTS) {
+		finished_updating = true;
+	}
+}
+
+TEST_CASE("Multiple concurrent updaters", "[interquery][.]") {
+	duckdb::unique_ptr<MaterializedQueryResult> result;
+	DuckDB db(nullptr);
+	Connection con(db);
+
+	// enable detailed profiling
+	con.Query("PRAGMA enable_profiling");
+	auto detailed_profiling_output = TestCreatePath("detailed_profiling_output");
+	con.Query("PRAGMA profiling_output='" + detailed_profiling_output + "'");
+	con.Query("PRAGMA profiling_mode = detailed");
+
+	finished_updating = false;
+	finished_threads = 0;
+	// initialize the database
+	con.Query("CREATE TABLE accounts(id INTEGER, money INTEGER)");
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TOTAL_ACCOUNTS; i++) {
+		con.Query("INSERT INTO accounts VALUES (" + to_string(i) + ", " +
+		          to_string(CONCURRENT_UPDATE_MONEY_PER_ACCOUNT) + ");");
+	}
+
+	bool correct[CONCURRENT_UPDATE_TOTAL_ACCOUNTS];
+	bool read_correct;
+	std::thread write_threads[CONCURRENT_UPDATE_TOTAL_ACCOUNTS];
+	// launch a thread for reading the table
+	thread read_thread(read_total_balance, &db, &read_correct);
+	// launch several threads for updating the table
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TOTAL_ACCOUNTS; i++) {
+		write_threads[i] = thread(write_random_numbers_to_account, &db, correct, i);
+	}
+	read_thread.join();
+	for (size_t i = 0; i < CONCURRENT_UPDATE_TOTAL_ACCOUNTS; i++) {
+		write_threads[i].join();
+		REQUIRE(correct[i]);
+	}
+}
+
+} // namespace test_concurrent_update