From ec6be4b056cad3813bb7e6e543bfb1e7b7f45d59 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Uzarski?= <mikolaj.uzarski@scylladb.com>
Date: Fri, 29 Nov 2024 13:43:23 +0100
Subject: [PATCH] argconv: make FFI implementation mutually exclusive

Disallow implementing two different APIs for specific type.
---
 scylla-rust-wrapper/src/argconv.rs      | 80 ++++++++++++++++++++++++-
 scylla-rust-wrapper/src/batch.rs        |  9 ++-
 scylla-rust-wrapper/src/cass_types.rs   |  4 +-
 scylla-rust-wrapper/src/cluster.rs      |  4 +-
 scylla-rust-wrapper/src/collection.rs   |  4 +-
 scylla-rust-wrapper/src/exec_profile.rs |  9 ++-
 scylla-rust-wrapper/src/future.rs       |  4 +-
 scylla-rust-wrapper/src/logging.rs      |  8 ++-
 scylla-rust-wrapper/src/metadata.rs     | 20 +++++--
 scylla-rust-wrapper/src/prepared.rs     |  4 +-
 scylla-rust-wrapper/src/query_error.rs  |  4 +-
 scylla-rust-wrapper/src/query_result.rs | 16 +++--
 scylla-rust-wrapper/src/retry_policy.rs |  6 +-
 scylla-rust-wrapper/src/session.rs      |  4 +-
 scylla-rust-wrapper/src/ssl.rs          |  6 +-
 scylla-rust-wrapper/src/statement.rs    |  4 +-
 scylla-rust-wrapper/src/tuple.rs        |  4 +-
 scylla-rust-wrapper/src/user_type.rs    |  4 +-
 scylla-rust-wrapper/src/uuid.rs         |  4 +-
 19 files changed, 166 insertions(+), 32 deletions(-)

diff --git a/scylla-rust-wrapper/src/argconv.rs b/scylla-rust-wrapper/src/argconv.rs
index 9a06ed09..95053588 100644
--- a/scylla-rust-wrapper/src/argconv.rs
+++ b/scylla-rust-wrapper/src/argconv.rs
@@ -343,12 +343,18 @@ impl<T: Sized> CassPtr<T, (Borrowed, Const)> {
     }
 }
 
+mod own_sealed {
+    pub trait ExclusiveSealed {}
+    pub trait SharedSealed {}
+    pub trait BorrowedSealed {}
+}
+
 /// Defines a pointer manipulation API for non-shared heap-allocated data.
 ///
 /// Implement this trait for types that are allocated by the driver via [`Box::new`],
 /// and then returned to the user as a pointer. The user is responsible for freeing
 /// the memory associated with the pointer using corresponding driver's API function.
-pub trait BoxFFI: Sized {
+pub trait BoxFFI: Sized + own_sealed::ExclusiveSealed {
     fn into_ptr<M: Mutability>(self: Box<Self>) -> CassExclusivePtr<Self, M> {
         CassExclusivePtr::from_box(self)
     }
@@ -382,7 +388,7 @@ pub trait BoxFFI: Sized {
 /// The data should be allocated via [`Arc::new`], and then returned to the user as a pointer.
 /// The user is responsible for freeing the memory associated
 /// with the pointer using corresponding driver's API function.
-pub trait ArcFFI: Sized {
+pub trait ArcFFI: Sized + own_sealed::SharedSealed {
     fn as_ptr(self: &Arc<Self>) -> CassSharedPtr<Self> {
         CassSharedPtr::from_ref(self)
     }
@@ -423,7 +429,7 @@ pub trait ArcFFI: Sized {
 /// For example: lifetime of CassRow is bound by the lifetime of CassResult.
 /// There is no API function that frees the CassRow. It should be automatically
 /// freed when user calls cass_result_free.
-pub trait RefFFI: Sized {
+pub trait RefFFI: Sized + own_sealed::BorrowedSealed {
     fn as_ptr(&self) -> CassBorrowedPtr<Self> {
         CassBorrowedPtr::from_ref(self)
     }
@@ -440,3 +446,71 @@ pub trait RefFFI: Sized {
         ptr.is_null()
     }
 }
+
+/// This trait should be implemented for types that are passed between
+/// C and Rust API. We currently distinguish 3 kinds of implementors,
+/// wrt. the ownership. The implementor should pick one of the 3 ownership
+/// kinds as the associated type:
+/// - [`OwnershipExclusive`]
+/// - [`OwnershipShared`]
+/// - [`OwnershipBorrowed`]
+#[allow(clippy::upper_case_acronyms)]
+pub trait FFI {
+    type Ownership;
+}
+
+/// Represents types with an exclusive ownership.
+///
+/// Use this associated type for implementors that require:
+/// - [`CassExclusivePtr`] manipulation via [`BoxFFI`]
+/// - exclusive ownership of the corresponding object
+/// - potential mutability of the corresponding object
+/// - manual memory freeing
+///
+/// C API user should be responsible for freeing associated memory manually
+/// via corresponding API call.
+///
+/// An example of such implementor would be [`CassCluster`](crate::cluster::CassCluster):
+/// - it is allocated on the heap via [`Box::new`]
+/// - user is the exclusive owner of the CassCluster object
+/// - there is no API to increase a reference count of CassCluster object
+/// - CassCluster is mutable via some API methods (`cass_cluster_set_*`)
+/// - user is responsible for freeing the associated memory (`cass_cluster_free`)
+pub struct OwnershipExclusive;
+impl<T> own_sealed::ExclusiveSealed for T where T: FFI<Ownership = OwnershipExclusive> {}
+impl<T> BoxFFI for T where T: FFI<Ownership = OwnershipExclusive> {}
+
+/// Represents types with a shared ownership.
+///
+/// Use this associated type for implementors that require:
+/// - [`CassSharedPtr`] manipulation via [`ArcFFI`]
+/// - shared ownership of the corresponding object
+/// - manual memory freeing
+///
+/// C API user should be responsible for freeing (decreasing reference count of)
+/// associated memory manually via corresponding API call.
+///
+/// An example of such implementor would be [`CassDataType`](crate::cass_types::CassDataType):
+/// - it is allocated on the heap via [`Arc::new`]
+/// - there are multiple owners of the shared CassDataType object
+/// - some API functions require to increase a reference count of the object
+/// - user is responsible for freeing (decreasing RC of) the associated memory (`cass_data_type_free`)
+pub struct OwnershipShared;
+impl<T> own_sealed::SharedSealed for T where T: FFI<Ownership = OwnershipShared> {}
+impl<T> ArcFFI for T where T: FFI<Ownership = OwnershipShared> {}
+
+/// Represents borrowed types.
+///
+/// Use this associated type for implementors that do not require any assumptions
+/// about the pointer type (apart from validity).
+/// The implementation will enable [`CassBorrowedPtr`] manipulation via [`RefFFI`]
+///
+/// C API user is not responsible for freeing associated memory manually. The memory
+/// should be freed automatically, when the owner is being dropped.
+///
+/// An example of such implementor would be [`CassRow`](crate::query_result::CassRow):
+/// - its lifetime is tied to the lifetime of CassResult
+/// - user only "borrows" the pointer - he is not responsible for freeing the memory
+pub struct OwnershipBorrowed;
+impl<T> own_sealed::BorrowedSealed for T where T: FFI<Ownership = OwnershipBorrowed> {}
+impl<T> RefFFI for T where T: FFI<Ownership = OwnershipBorrowed> {}
diff --git a/scylla-rust-wrapper/src/batch.rs b/scylla-rust-wrapper/src/batch.rs
index a6bae761..4f90694f 100644
--- a/scylla-rust-wrapper/src/batch.rs
+++ b/scylla-rust-wrapper/src/batch.rs
@@ -1,4 +1,7 @@
-use crate::argconv::{ArcFFI, BoxFFI, CassExclusiveConstPtr, CassExclusiveMutPtr, CassSharedPtr};
+use crate::argconv::{
+    ArcFFI, BoxFFI, CassExclusiveConstPtr, CassExclusiveMutPtr, CassSharedPtr, OwnershipExclusive,
+    FFI,
+};
 use crate::cass_error::CassError;
 use crate::cass_types::CassConsistency;
 use crate::cass_types::{make_batch_type, CassBatchType};
@@ -19,7 +22,9 @@ pub struct CassBatch {
     pub(crate) exec_profile: Option<PerStatementExecProfile>,
 }
 
-impl BoxFFI for CassBatch {}
+impl FFI for CassBatch {
+    type Ownership = OwnershipExclusive;
+}
 
 #[derive(Clone)]
 pub struct CassBatchState {
diff --git a/scylla-rust-wrapper/src/cass_types.rs b/scylla-rust-wrapper/src/cass_types.rs
index 5a9a7eae..09414cfd 100644
--- a/scylla-rust-wrapper/src/cass_types.rs
+++ b/scylla-rust-wrapper/src/cass_types.rs
@@ -175,7 +175,9 @@ pub enum CassDataTypeInner {
     Custom(String),
 }
 
-impl ArcFFI for CassDataType {}
+impl FFI for CassDataType {
+    type Ownership = OwnershipShared;
+}
 
 impl CassDataTypeInner {
     /// Checks for equality during typechecks.
diff --git a/scylla-rust-wrapper/src/cluster.rs b/scylla-rust-wrapper/src/cluster.rs
index 11621418..29ea3dbb 100644
--- a/scylla-rust-wrapper/src/cluster.rs
+++ b/scylla-rust-wrapper/src/cluster.rs
@@ -165,7 +165,9 @@ impl CassCluster {
     }
 }
 
-impl BoxFFI for CassCluster {}
+impl FFI for CassCluster {
+    type Ownership = OwnershipExclusive;
+}
 
 pub struct CassCustomPayload;
 
diff --git a/scylla-rust-wrapper/src/collection.rs b/scylla-rust-wrapper/src/collection.rs
index 527c219f..27bd16ea 100644
--- a/scylla-rust-wrapper/src/collection.rs
+++ b/scylla-rust-wrapper/src/collection.rs
@@ -36,7 +36,9 @@ pub struct CassCollection {
     pub items: Vec<CassCqlValue>,
 }
 
-impl BoxFFI for CassCollection {}
+impl FFI for CassCollection {
+    type Ownership = OwnershipExclusive;
+}
 
 impl CassCollection {
     fn typecheck_on_append(&self, value: &Option<CassCqlValue>) -> CassError {
diff --git a/scylla-rust-wrapper/src/exec_profile.rs b/scylla-rust-wrapper/src/exec_profile.rs
index 0d866751..f1ef4014 100644
--- a/scylla-rust-wrapper/src/exec_profile.rs
+++ b/scylla-rust-wrapper/src/exec_profile.rs
@@ -13,7 +13,10 @@ use scylla::retry_policy::RetryPolicy;
 use scylla::speculative_execution::SimpleSpeculativeExecutionPolicy;
 use scylla::statement::Consistency;
 
-use crate::argconv::{ptr_to_cstr_n, strlen, ArcFFI, BoxFFI, CassExclusiveMutPtr, CassSharedPtr};
+use crate::argconv::{
+    ptr_to_cstr_n, strlen, ArcFFI, BoxFFI, CassExclusiveMutPtr, CassSharedPtr, OwnershipExclusive,
+    FFI,
+};
 use crate::batch::CassBatch;
 use crate::cass_error::CassError;
 use crate::cass_types::CassConsistency;
@@ -37,7 +40,9 @@ pub struct CassExecProfile {
     load_balancing_config: LoadBalancingConfig,
 }
 
-impl BoxFFI for CassExecProfile {}
+impl FFI for CassExecProfile {
+    type Ownership = OwnershipExclusive;
+}
 
 impl CassExecProfile {
     fn new() -> Self {
diff --git a/scylla-rust-wrapper/src/future.rs b/scylla-rust-wrapper/src/future.rs
index 42fcc511..5cfc71f4 100644
--- a/scylla-rust-wrapper/src/future.rs
+++ b/scylla-rust-wrapper/src/future.rs
@@ -60,7 +60,9 @@ pub struct CassFuture {
     wait_for_value: Condvar,
 }
 
-impl ArcFFI for CassFuture {}
+impl FFI for CassFuture {
+    type Ownership = OwnershipShared;
+}
 
 /// An error that can appear during `cass_future_wait_timed`.
 enum FutureError {
diff --git a/scylla-rust-wrapper/src/logging.rs b/scylla-rust-wrapper/src/logging.rs
index 0aca2b53..70fe468c 100644
--- a/scylla-rust-wrapper/src/logging.rs
+++ b/scylla-rust-wrapper/src/logging.rs
@@ -1,4 +1,6 @@
-use crate::argconv::{arr_to_cstr, ptr_to_cstr, str_to_arr, CassBorrowedPtr, RefFFI};
+use crate::argconv::{
+    arr_to_cstr, ptr_to_cstr, str_to_arr, CassBorrowedPtr, OwnershipBorrowed, RefFFI, FFI,
+};
 use crate::cass_log_types::{CassLogLevel, CassLogMessage};
 use crate::types::size_t;
 use crate::LOGGER;
@@ -14,7 +16,9 @@ use tracing_subscriber::layer::Context;
 use tracing_subscriber::prelude::*;
 use tracing_subscriber::Layer;
 
-impl RefFFI for CassLogMessage {}
+impl FFI for CassLogMessage {
+    type Ownership = OwnershipBorrowed;
+}
 
 pub type CassLogCallback =
     Option<unsafe extern "C" fn(message: CassBorrowedPtr<CassLogMessage>, data: *mut c_void)>;
diff --git a/scylla-rust-wrapper/src/metadata.rs b/scylla-rust-wrapper/src/metadata.rs
index 2c515486..b81a82b6 100644
--- a/scylla-rust-wrapper/src/metadata.rs
+++ b/scylla-rust-wrapper/src/metadata.rs
@@ -13,7 +13,9 @@ pub struct CassSchemaMeta {
     pub keyspaces: HashMap<String, CassKeyspaceMeta>,
 }
 
-impl BoxFFI for CassSchemaMeta {}
+impl FFI for CassSchemaMeta {
+    type Ownership = OwnershipExclusive;
+}
 
 pub struct CassKeyspaceMeta {
     pub name: String,
@@ -25,7 +27,9 @@ pub struct CassKeyspaceMeta {
 }
 
 // Owned by CassSchemaMeta
-impl RefFFI for CassKeyspaceMeta {}
+impl FFI for CassKeyspaceMeta {
+    type Ownership = OwnershipBorrowed;
+}
 
 pub struct CassTableMeta {
     pub name: String,
@@ -38,7 +42,9 @@ pub struct CassTableMeta {
 // Either:
 // - owned by CassMaterializedViewMeta - won't be given to user
 // - Owned by CassKeyspaceMeta (in Arc), referenced (Weak) by CassMaterializedViewMeta
-impl RefFFI for CassTableMeta {}
+impl FFI for CassTableMeta {
+    type Ownership = OwnershipBorrowed;
+}
 
 pub struct CassMaterializedViewMeta {
     pub name: String,
@@ -47,7 +53,9 @@ pub struct CassMaterializedViewMeta {
 }
 
 // Shared ownership by CassKeyspaceMeta and CassTableMeta
-impl RefFFI for CassMaterializedViewMeta {}
+impl FFI for CassMaterializedViewMeta {
+    type Ownership = OwnershipBorrowed;
+}
 
 pub struct CassColumnMeta {
     pub name: String,
@@ -56,7 +64,9 @@ pub struct CassColumnMeta {
 }
 
 // Owned by CassTableMeta
-impl RefFFI for CassColumnMeta {}
+impl FFI for CassColumnMeta {
+    type Ownership = OwnershipBorrowed;
+}
 
 pub unsafe fn create_table_metadata(
     keyspace_name: &str,
diff --git a/scylla-rust-wrapper/src/prepared.rs b/scylla-rust-wrapper/src/prepared.rs
index 9ea0515a..2f923ad2 100644
--- a/scylla-rust-wrapper/src/prepared.rs
+++ b/scylla-rust-wrapper/src/prepared.rs
@@ -72,7 +72,9 @@ impl CassPrepared {
     }
 }
 
-impl ArcFFI for CassPrepared {}
+impl FFI for CassPrepared {
+    type Ownership = OwnershipShared;
+}
 
 #[no_mangle]
 pub unsafe extern "C" fn cass_prepared_free(prepared_raw: CassSharedPtr<CassPrepared>) {
diff --git a/scylla-rust-wrapper/src/query_error.rs b/scylla-rust-wrapper/src/query_error.rs
index a8898014..db11315a 100644
--- a/scylla-rust-wrapper/src/query_error.rs
+++ b/scylla-rust-wrapper/src/query_error.rs
@@ -19,7 +19,9 @@ pub enum CassErrorResult {
     Deserialization(#[from] DeserializationError),
 }
 
-impl ArcFFI for CassErrorResult {}
+impl FFI for CassErrorResult {
+    type Ownership = OwnershipShared;
+}
 
 impl From<Consistency> for CassConsistency {
     fn from(c: Consistency) -> CassConsistency {
diff --git a/scylla-rust-wrapper/src/query_result.rs b/scylla-rust-wrapper/src/query_result.rs
index e21889a3..78530274 100644
--- a/scylla-rust-wrapper/src/query_result.rs
+++ b/scylla-rust-wrapper/src/query_result.rs
@@ -95,7 +95,9 @@ impl CassResult {
     }
 }
 
-impl ArcFFI for CassResult {}
+impl FFI for CassResult {
+    type Ownership = OwnershipShared;
+}
 
 #[derive(Debug)]
 pub struct CassResultMetadata {
@@ -149,7 +151,9 @@ pub struct CassRow {
     pub result_metadata: Arc<CassResultMetadata>,
 }
 
-impl RefFFI for CassRow {}
+impl FFI for CassRow {
+    type Ownership = OwnershipBorrowed;
+}
 
 pub fn create_cass_rows_from_rows(
     rows: Vec<Row>,
@@ -185,7 +189,9 @@ pub struct CassValue {
     pub value_type: Arc<CassDataType>,
 }
 
-impl RefFFI for CassValue {}
+impl FFI for CassValue {
+    type Ownership = OwnershipBorrowed;
+}
 
 fn create_cass_row_columns(row: Row, metadata: &Arc<CassResultMetadata>) -> Vec<CassValue> {
     row.columns
@@ -367,7 +373,9 @@ pub enum CassIterator {
     CassViewMetaIterator(CassViewMetaIterator),
 }
 
-impl BoxFFI for CassIterator {}
+impl FFI for CassIterator {
+    type Ownership = OwnershipExclusive;
+}
 
 #[no_mangle]
 pub unsafe extern "C" fn cass_iterator_free(iterator: CassExclusiveMutPtr<CassIterator>) {
diff --git a/scylla-rust-wrapper/src/retry_policy.rs b/scylla-rust-wrapper/src/retry_policy.rs
index 210c76d5..0fd451c9 100644
--- a/scylla-rust-wrapper/src/retry_policy.rs
+++ b/scylla-rust-wrapper/src/retry_policy.rs
@@ -2,7 +2,7 @@ use scylla::retry_policy::{DefaultRetryPolicy, FallthroughRetryPolicy};
 use scylla::transport::downgrading_consistency_retry_policy::DowngradingConsistencyRetryPolicy;
 use std::sync::Arc;
 
-use crate::argconv::{ArcFFI, CassSharedPtr};
+use crate::argconv::{ArcFFI, CassSharedPtr, OwnershipShared, FFI};
 
 pub enum RetryPolicy {
     DefaultRetryPolicy(Arc<DefaultRetryPolicy>),
@@ -12,7 +12,9 @@ pub enum RetryPolicy {
 
 pub type CassRetryPolicy = RetryPolicy;
 
-impl ArcFFI for CassRetryPolicy {}
+impl FFI for CassRetryPolicy {
+    type Ownership = OwnershipShared;
+}
 
 #[no_mangle]
 pub extern "C" fn cass_retry_policy_default_new() -> CassSharedPtr<CassRetryPolicy> {
diff --git a/scylla-rust-wrapper/src/session.rs b/scylla-rust-wrapper/src/session.rs
index 62b06dc5..8d003751 100644
--- a/scylla-rust-wrapper/src/session.rs
+++ b/scylla-rust-wrapper/src/session.rs
@@ -139,7 +139,9 @@ impl CassSessionInner {
 
 pub type CassSession = RwLock<Option<CassSessionInner>>;
 
-impl ArcFFI for CassSession {}
+impl FFI for CassSession {
+    type Ownership = OwnershipShared;
+}
 
 #[no_mangle]
 pub unsafe extern "C" fn cass_session_new() -> CassSharedPtr<CassSession> {
diff --git a/scylla-rust-wrapper/src/ssl.rs b/scylla-rust-wrapper/src/ssl.rs
index 70cf992d..7b285fb3 100644
--- a/scylla-rust-wrapper/src/ssl.rs
+++ b/scylla-rust-wrapper/src/ssl.rs
@@ -1,5 +1,7 @@
 use crate::argconv::ArcFFI;
 use crate::argconv::CassSharedPtr;
+use crate::argconv::OwnershipShared;
+use crate::argconv::FFI;
 use crate::cass_error::CassError;
 use crate::types::size_t;
 use libc::{c_int, strlen};
@@ -20,7 +22,9 @@ pub struct CassSsl {
     pub(crate) trusted_store: *mut X509_STORE,
 }
 
-impl ArcFFI for CassSsl {}
+impl FFI for CassSsl {
+    type Ownership = OwnershipShared;
+}
 
 pub const CASS_SSL_VERIFY_NONE: i32 = 0x00;
 pub const CASS_SSL_VERIFY_PEER_CERT: i32 = 0x01;
diff --git a/scylla-rust-wrapper/src/statement.rs b/scylla-rust-wrapper/src/statement.rs
index e1e354ec..c19c5af2 100644
--- a/scylla-rust-wrapper/src/statement.rs
+++ b/scylla-rust-wrapper/src/statement.rs
@@ -42,7 +42,9 @@ pub struct CassStatement {
     pub(crate) exec_profile: Option<PerStatementExecProfile>,
 }
 
-impl BoxFFI for CassStatement {}
+impl FFI for CassStatement {
+    type Ownership = OwnershipExclusive;
+}
 
 impl CassStatement {
     fn bind_cql_value(&mut self, index: usize, value: Option<CassCqlValue>) -> CassError {
diff --git a/scylla-rust-wrapper/src/tuple.rs b/scylla-rust-wrapper/src/tuple.rs
index 3d385a32..1d83da7f 100644
--- a/scylla-rust-wrapper/src/tuple.rs
+++ b/scylla-rust-wrapper/src/tuple.rs
@@ -17,7 +17,9 @@ pub struct CassTuple {
     pub items: Vec<Option<CassCqlValue>>,
 }
 
-impl BoxFFI for CassTuple {}
+impl FFI for CassTuple {
+    type Ownership = OwnershipExclusive;
+}
 
 impl CassTuple {
     fn get_types(&self) -> Option<&Vec<Arc<CassDataType>>> {
diff --git a/scylla-rust-wrapper/src/user_type.rs b/scylla-rust-wrapper/src/user_type.rs
index 4b160109..b4189384 100644
--- a/scylla-rust-wrapper/src/user_type.rs
+++ b/scylla-rust-wrapper/src/user_type.rs
@@ -14,7 +14,9 @@ pub struct CassUserType {
     pub field_values: Vec<Option<CassCqlValue>>,
 }
 
-impl BoxFFI for CassUserType {}
+impl FFI for CassUserType {
+    type Ownership = OwnershipExclusive;
+}
 
 impl CassUserType {
     fn set_field_by_index(&mut self, index: usize, value: Option<CassCqlValue>) -> CassError {
diff --git a/scylla-rust-wrapper/src/uuid.rs b/scylla-rust-wrapper/src/uuid.rs
index 29ec4772..e4d57be1 100644
--- a/scylla-rust-wrapper/src/uuid.rs
+++ b/scylla-rust-wrapper/src/uuid.rs
@@ -17,7 +17,9 @@ pub struct CassUuidGen {
     pub last_timestamp: AtomicU64,
 }
 
-impl BoxFFI for CassUuidGen {}
+impl FFI for CassUuidGen {
+    type Ownership = OwnershipExclusive;
+}
 
 // Implementation directly ported from Cpp Driver implementation: