From d727e7648ecd23010d2561eed064ea83f5d33b9a Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew.noyes@snowflake.com>
Date: Mon, 28 Mar 2022 14:33:59 -0700
Subject: [PATCH] Fix a few memory issues found by ASAN

---
 bindings/c/test/unit/fdb_api.cpp |  8 +++++++-
 bindings/c/test/unit/fdb_api.hpp |  7 ++++++-
 fdbcli/TenantCommands.actor.cpp  | 18 +++++++++++++-----
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/bindings/c/test/unit/fdb_api.cpp b/bindings/c/test/unit/fdb_api.cpp
index 4fc715dbc51..b26d7bdf82e 100644
--- a/bindings/c/test/unit/fdb_api.cpp
+++ b/bindings/c/test/unit/fdb_api.cpp
@@ -138,6 +138,12 @@ Tenant::Tenant(FDBDatabase* db, const uint8_t* name, int name_length) {
 	}
 }
 
+Tenant::~Tenant() {
+	if (tenant != nullptr) {
+		fdb_tenant_destroy(tenant);
+	}
+}
+
 // Transaction
 Transaction::Transaction(FDBDatabase* db) {
 	if (fdb_error_t err = fdb_database_create_transaction(db, &tr_)) {
@@ -146,7 +152,7 @@ Transaction::Transaction(FDBDatabase* db) {
 	}
 }
 
-Transaction::Transaction(Tenant tenant) {
+Transaction::Transaction(Tenant& tenant) {
 	if (fdb_error_t err = fdb_tenant_create_transaction(tenant.tenant, &tr_)) {
 		std::cerr << fdb_get_error(err) << std::endl;
 		std::abort();
diff --git a/bindings/c/test/unit/fdb_api.hpp b/bindings/c/test/unit/fdb_api.hpp
index 5653d6e7cb5..fcf1c7e5ca1 100644
--- a/bindings/c/test/unit/fdb_api.hpp
+++ b/bindings/c/test/unit/fdb_api.hpp
@@ -206,6 +206,11 @@ class Database final {
 class Tenant final {
 public:
 	Tenant(FDBDatabase* db, const uint8_t* name, int name_length);
+	~Tenant();
+	Tenant(const Tenant&) = delete;
+	Tenant& operator=(const Tenant&) = delete;
+	Tenant(Tenant&&) = delete;
+	Tenant& operator=(Tenant&&) = delete;
 
 private:
 	friend class Transaction;
@@ -219,7 +224,7 @@ class Transaction final {
 public:
 	// Given an FDBDatabase, initializes a new transaction.
 	Transaction(FDBDatabase* db);
-	Transaction(Tenant tenant);
+	Transaction(Tenant& tenant);
 	~Transaction();
 
 	// Wrapper around fdb_transaction_reset.
diff --git a/fdbcli/TenantCommands.actor.cpp b/fdbcli/TenantCommands.actor.cpp
index c03bb17c882..6660893a1f6 100644
--- a/fdbcli/TenantCommands.actor.cpp
+++ b/fdbcli/TenantCommands.actor.cpp
@@ -51,7 +51,9 @@ ACTOR Future<bool> createTenantCommandActor(Reference<IDatabase> db, std::vector
 		tr->setOption(FDBTransactionOptions::SPECIAL_KEY_SPACE_ENABLE_WRITES);
 		try {
 			if (!doneExistenceCheck) {
-				Optional<Value> existingTenant = wait(safeThreadFutureToFuture(tr->get(tenantNameKey)));
+				// Hold the reference to the standalone's memory
+				state ThreadFuture<Optional<Value>> existingTenantFuture = tr->get(tenantNameKey);
+				Optional<Value> existingTenant = wait(safeThreadFutureToFuture(existingTenantFuture));
 				if (existingTenant.present()) {
 					throw tenant_already_exists();
 				}
@@ -96,7 +98,9 @@ ACTOR Future<bool> deleteTenantCommandActor(Reference<IDatabase> db, std::vector
 		tr->setOption(FDBTransactionOptions::SPECIAL_KEY_SPACE_ENABLE_WRITES);
 		try {
 			if (!doneExistenceCheck) {
-				Optional<Value> existingTenant = wait(safeThreadFutureToFuture(tr->get(tenantNameKey)));
+				// Hold the reference to the standalone's memory
+				state ThreadFuture<Optional<Value>> existingTenantFuture = tr->get(tenantNameKey);
+				Optional<Value> existingTenant = wait(safeThreadFutureToFuture(existingTenantFuture));
 				if (!existingTenant.present()) {
 					throw tenant_not_found();
 				}
@@ -163,8 +167,10 @@ ACTOR Future<bool> listTenantsCommandActor(Reference<IDatabase> db, std::vector<
 
 	loop {
 		try {
-			RangeResult tenants = wait(safeThreadFutureToFuture(
-			    tr->getRange(firstGreaterOrEqual(beginTenantKey), firstGreaterOrEqual(endTenantKey), limit)));
+			// Hold the reference to the standalone's memory
+			state ThreadFuture<RangeResult> kvsFuture =
+			    tr->getRange(firstGreaterOrEqual(beginTenantKey), firstGreaterOrEqual(endTenantKey), limit);
+			RangeResult tenants = wait(safeThreadFutureToFuture(kvsFuture));
 
 			if (tenants.empty()) {
 				if (tokens.size() == 1) {
@@ -213,7 +219,9 @@ ACTOR Future<bool> getTenantCommandActor(Reference<IDatabase> db, std::vector<St
 
 	loop {
 		try {
-			Optional<Value> tenant = wait(safeThreadFutureToFuture(tr->get(tenantNameKey)));
+			// Hold the reference to the standalone's memory
+			state ThreadFuture<Optional<Value>> tenantFuture = tr->get(tenantNameKey);
+			Optional<Value> tenant = wait(safeThreadFutureToFuture(tenantFuture));
 			if (!tenant.present()) {
 				throw tenant_not_found();
 			}