Re: [PR] Implement `copy_if_not_exist` for `AmazonS3` using DynamoDB (#4880) [arrow-rs]

["tustvold (via GitHub)" <gi...@apache.org>]

tustvold commented on code in PR #4918:
URL: https://github.com/apache/arrow-rs/pull/4918#discussion_r1435276250


##########
object_store/src/aws/dynamo.rs:
##########
@@ -0,0 +1,567 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! A DynamoDB based lock system
+
+use std::collections::HashMap;
+use std::time::{Duration, Instant};
+
+use chrono::Utc;
+use reqwest::{Response, StatusCode};
+use serde::ser::SerializeMap;
+use serde::{Deserialize, Serialize, Serializer};
+
+use crate::aws::client::S3Client;
+use crate::aws::credential::CredentialExt;
+use crate::aws::AwsCredential;
+use crate::client::get::GetClientExt;
+use crate::client::retry::Error as RetryError;
+use crate::client::retry::RetryExt;
+use crate::path::Path;
+use crate::{Error, GetOptions, Result};
+
+/// The exception returned by DynamoDB on conflict
+const CONFLICT: &str = "ConditionalCheckFailedException";
+
+const STORE: &str = "DynamoDB";
+
+/// A DynamoDB-based commit protocol, used to provide conditional write support for S3
+///
+/// ## Limitations
+///
+/// Only conditional operations, e.g. `copy_if_not_exists` will be synchronized, and can
+/// therefore race with non-conditional operations, e.g. `put`, `copy`, `delete`, or
+/// conditional operations performed by writers not configured to synchronize with DynamoDB.
+///
+/// Workloads making use of this mechanism **must** ensure:
+///
+/// * Conditional and non-conditional operations are not performed on the same paths
+/// * Conditional operations are only performed via similarly configured clients
+///
+/// Additionally as the locking mechanism relies on timeouts to detect stale locks,
+/// performance will be poor for systems that frequently delete and then create
+/// objects at the same path, instead being optimised for systems that primarily create
+/// files with paths never used before, or perform conditional updates to existing files
+///
+/// ## Commit Protocol
+///
+/// The DynamoDB schema is as follows:
+///
+/// * A string hash key named `"key"`
+/// * A numeric [TTL] attribute named `"ttl"`
+/// * A numeric attribute named `"generation"`
+/// * A numeric attribute named `"timeout"`
+///
+/// To perform a conditional operation on an object with a given `path` and `etag` (if exists),
+/// the commit protocol is as follows:
+///
+/// 1. Perform HEAD request on `path` and error on precondition mismatch
+/// 2. Create record in DynamoDB with key `{path}#{etag}` with the configured timeout
+///     1. On Success: Perform operation with the configured timeout
+///     2. On Conflict:
+///         1. Periodically re-perform HEAD request on `path` and error on precondition mismatch
+///         2. If `timeout * max_skew_rate` passed, replace the record incrementing the `"generation"`
+///             1. On Success: GOTO 2.1
+///             2. On Conflict: GOTO 2.2
+///
+/// Provided no writer modifies an object with a given `path` and `etag` without first adding a
+/// corresponding record to DynamoDB, we are guaranteed that only one writer will ever commit.
+///
+/// This is inspired by the [DynamoDB Lock Client] but simplified for the more limited
+/// requirements of synchronizing object storage. The major changes are:
+///
+/// * Uses a monotonic generation count instead of a UUID rvn, as this is:
+///     * Cheaper to generate, serialize and compare
+///     * Cannot collide
+///     * More human readable / interpretable
+/// * Relies on [TTL] to eventually clean up old locks
+///
+/// It also draws inspiration from the DeltaLake [S3 Multi-Cluster] commit protocol, but
+/// generalised to not make assumptions about the workload and not rely on first writing
+/// to a temporary path.
+///
+/// [TTL]: https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/howitworks-ttl.html
+/// [DynamoDB Lock Client]: https://aws.amazon.com/blogs/database/building-distributed-locks-with-the-dynamodb-lock-client/
+/// [S3 Multi-Cluster]: https://docs.google.com/document/d/1Gs4ZsTH19lMxth4BSdwlWjUNR-XhKHicDvBjd2RqNd8/edit#heading=h.mjjuxw9mcz9h
+#[derive(Debug, Clone)]
+pub struct DynamoCommit {
+    table_name: String,
+    /// The number of milliseconds a lease is valid for
+    timeout: u64,
+    /// The maximum clock skew rate tolerated by the system
+    max_clock_skew_rate: u32,
+    /// The length of time a record will be retained in DynamoDB before being cleaned up
+    ///
+    /// This is purely an optimisation to avoid indefinite growth of the DynamoDB table
+    /// and does not impact how long clients may wait to acquire a lock
+    ttl: Duration,
+    /// The backoff duration before retesting a condition
+    test_interval: Duration,
+}
+
+impl DynamoCommit {
+    /// Create a new [`DynamoCommit`] with a given table name
+    pub fn new(table_name: String) -> Self {
+        Self {
+            table_name,
+            timeout: 20_000,
+            max_clock_skew_rate: 3,
+            ttl: Duration::from_secs(60 * 60),
+            test_interval: Duration::from_millis(100),
+        }
+    }
+
+    /// Overrides the lock timeout.
+    ///
+    /// A longer lock timeout reduces the probability of spurious commit failures and multi-writer
+    /// races, but will increase the time that writers must wait to reclaim a lock lost. The
+    /// default value of 20 seconds should be appropriate for must use-cases.
+    pub fn with_timeout(mut self, millis: u64) -> Self {
+        self.timeout = millis;
+        self
+    }
+
+    /// The maximum clock skew rate tolerated by the system.
+    ///
+    /// An environment in which the clock on the fastest node ticks twice as fast as the slowest
+    /// node, would have a clock skew rate of 2. The default value of 3 should be appropriate
+    /// for most environments.
+    pub fn with_max_clock_skew_rate(mut self, rate: u32) -> Self {
+        self.max_clock_skew_rate = rate;
+        self
+    }
+
+    /// The length of time a record should be retained in DynamoDB before being cleaned up
+    ///
+    /// This should be significantly larger than the configured lock timeout, with the default
+    /// value of 1 hour appropriate for most use-cases.
+    pub fn with_ttl(mut self, ttl: Duration) -> Self {
+        self.ttl = ttl;
+        self
+    }
+
+    /// Returns the name of the DynamoDB table.
+    pub(crate) fn table_name(&self) -> &str {
+        &self.table_name
+    }
+
+    pub(crate) async fn copy_if_not_exists(
+        &self,
+        client: &S3Client,
+        from: &Path,
+        to: &Path,
+    ) -> Result<()> {
+        check_not_exists(client, to).await?;
+
+        let mut previous_lease = None;
+
+        loop {
+            let existing = previous_lease.as_ref();
+            match self.try_lock(client, to.as_ref(), existing).await? {
+                TryLockResult::Ok(lease) => {
+                    let fut = client.copy_request(from, to).send();
+                    let expiry = lease.acquire + lease.timeout;
+                    return match tokio::time::timeout_at(expiry.into(), fut).await {
+                        Ok(Ok(_)) => Ok(()),
+                        Ok(Err(e)) => Err(e),
+                        Err(_) => Err(Error::Generic {
+                            store: "DynamoDB",
+                            source: format!(
+                                "Failed to perform copy operation in {} milliseconds",
+                                self.timeout
+                            )
+                            .into(),
+                        }),
+                    };
+                }
+                TryLockResult::Conflict(conflict) => {
+                    let mut interval = tokio::time::interval(self.test_interval);
+                    let expiry = conflict.timeout * self.max_clock_skew_rate;
+                    loop {
+                        interval.tick().await;
+                        check_not_exists(client, to).await?;
+                        if conflict.acquire.elapsed() > expiry {
+                            previous_lease = Some(conflict);
+                            break;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    /// Retrieve a lock, returning an error if it doesn't exist
+    async fn get_lock(&self, s3: &S3Client, key: &str) -> Result<Lease> {
+        let key_attributes = [("key", AttributeValue::String(key))];
+        let req = GetItem {
+            table_name: &self.table_name,
+            key: Map(&key_attributes),
+        };
+        let credential = s3.config.get_credential().await?;
+
+        let resp = self
+            .request(s3, credential.as_deref(), "DynamoDB_20120810.GetItem", req)
+            .await
+            .map_err(|e| e.error(STORE, key.to_string()))?;
+
+        let body = resp.bytes().await.map_err(|e| Error::Generic {
+            store: STORE,
+            source: Box::new(e),
+        })?;
+
+        let response: GetItemResponse<'_> =
+            serde_json::from_slice(body.as_ref()).map_err(|e| Error::Generic {
+                store: STORE,
+                source: Box::new(e),
+            })?;
+
+        extract_lease(&response.item).ok_or_else(|| Error::NotFound {
+            path: key.into(),
+            source: "DynamoDB GetItem returned no items".to_string().into(),
+        })
+    }
+
+    /// Attempt to acquire a lock, reclaiming an existing lease if provided
+    async fn try_lock(
+        &self,
+        s3: &S3Client,
+        key: &str,
+        existing: Option<&Lease>,
+    ) -> Result<TryLockResult> {
+        let attributes;
+        let (next_gen, condition_expression, expression_attribute_values) = match existing {
+            None => (0_u64, "attribute_not_exists(#pk)", Map(&[])),
+            Some(existing) => {
+                attributes = [(":g", AttributeValue::Number(existing.generation))];
+                (
+                    existing.generation.checked_add(1).unwrap(),
+                    "attribute_exists(#pk) AND generation = :g",
+                    Map(attributes.as_slice()),
+                )
+            }
+        };
+
+        let ttl = (Utc::now() + self.ttl).timestamp();
+        let items = [
+            ("key", AttributeValue::String(key)),
+            ("generation", AttributeValue::Number(next_gen)),
+            ("timeout", AttributeValue::Number(self.timeout)),
+            ("ttl", AttributeValue::Number(ttl as _)),

Review Comment:
   It is documented here - https://github.com/apache/arrow-rs/pull/4918/files#diff-5a8585cdad662383c911426918d89de6d4b7ed13b975767c96ecac30d6e0c227R62



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org