feat: optimistic concurrency control loop for IAM (#1162)

Adds functions to DatabaseAdminClient and InstanceAdminClient that run a
optimistic concurrency control (OCC) loop for IAM updates. That makes
the customer code a bit easier to write and less error prone. This is
analogous to the Commit() loops for the spanner::Client.

Author: coryan

google/cloud/spanner/CMakeLists.txt

@@ -117,6 +117,7 @@ add_library(
     database_admin_connection.h
     date.cc
     date.h
+    iam_updater.h
     instance.cc
     instance.h
     instance_admin_client.cc

google/cloud/spanner/database_admin_client.cc

@@ -64,6 +64,54 @@ StatusOr<google::iam::v1::Policy> DatabaseAdminClient::SetIamPolicy(
   return conn_->SetIamPolicy({std::move(db), std::move(policy)});
 }
 
+StatusOr<google::iam::v1::Policy> DatabaseAdminClient::SetIamPolicy(
+    Database const& db, IamUpdater const& updater) {
+  auto const rerun_maximum_duration = std::chrono::minutes(15);
+  auto default_rerun_policy =
+      LimitedTimeTransactionRerunPolicy(rerun_maximum_duration).clone();
+
+  auto const backoff_initial_delay = std::chrono::milliseconds(1000);
+  auto const backoff_maximum_delay = std::chrono::minutes(5);
+  auto const backoff_scaling = 2.0;
+  auto default_backoff_policy =
+      ExponentialBackoffPolicy(backoff_initial_delay, backoff_maximum_delay,
+                               backoff_scaling)
+          .clone();
+
+  return SetIamPolicy(db, updater, std::move(default_rerun_policy),
+                      std::move(default_backoff_policy));
+}
+
+StatusOr<google::iam::v1::Policy> DatabaseAdminClient::SetIamPolicy(
+    Database const& db, IamUpdater const& updater,
+    std::unique_ptr<TransactionRerunPolicy> rerun_policy,
+    std::unique_ptr<BackoffPolicy> backoff_policy) {
+  using RerunnablePolicy = internal::SafeTransactionRerun;
+
+  Status last_status;
+  do {
+    auto current_policy = GetIamPolicy(db);
+    if (!current_policy) {
+      last_status = std::move(current_policy).status();
+    } else {
+      auto etag = current_policy->etag();
+      auto desired = updater(*current_policy);
+      if (!desired.has_value()) {
+        return current_policy;
+      }
+      desired->set_etag(std::move(etag));
+      auto result = SetIamPolicy(db, *std::move(desired));
+      if (RerunnablePolicy::IsOk(result.status())) {
+        return result;
+      }
+      last_status = std::move(result).status();
+    }
+    if (!rerun_policy->OnFailure(last_status)) break;
+    std::this_thread::sleep_for(backoff_policy->OnCompletion());
+  } while (!rerun_policy->IsExhausted());
+  return last_status;
+}
+
 StatusOr<google::iam::v1::TestIamPermissionsResponse>
 DatabaseAdminClient::TestIamPermissions(Database db,
                                         std::vector<std::string> permissions) {

google/cloud/spanner/database_admin_client.h

@@ -18,6 +18,7 @@
 #include "google/cloud/spanner/connection_options.h"
 #include "google/cloud/spanner/database.h"
 #include "google/cloud/spanner/database_admin_connection.h"
+#include "google/cloud/spanner/iam_updater.h"
 #include "google/cloud/spanner/instance.h"
 #include "google/cloud/future.h"
 #include "google/cloud/status_or.h"
@@ -231,9 +232,6 @@ class DatabaseAdminClient {
    * Therefore, the underlying RPCs are only retried if the field is set, and
    * the function returns the first RPC error in any other case.
    *
-   * @par Example
-   * @snippet samples.cc add-database-reader-on-database
-   *
    * @see The [Cloud Spanner
    *     documentation](https://cloud.google.com/spanner/docs/iam) for a
    *     description of the roles and permissions supported by Cloud Spanner.
@@ -244,6 +242,50 @@ class DatabaseAdminClient {
   StatusOr<google::iam::v1::Policy> SetIamPolicy(
       Database db, google::iam::v1::Policy policy);
 
+  /**
+   * Updates the IAM policy for an instance using an optimistic concurrency
+   * control loop.
+   *
+   * This function repeatedly reads the current IAM policy in @p db, and then
+   * calls the @p updater with the this policy. The @p updater returns an empty
+   * optional if no changes are required, or it returns the new desired value
+   * for the IAM policy. This function then updates the policy.
+   *
+   * Updating an IAM policy can fail with retryable errors or can be aborted
+   * because there were simultaneous changes the to IAM policy. In these cases
+   * this function reruns the loop until it succeeds.
+   *
+   * The function returns the final IAM policy, or an error if the rerun policy
+   * for the underlying connection has expired.
+   *
+   * @par Idempotency
+   * This function always sets the `etag` field on the policy, so the underlying
+   * RPCs are retried automatically.
+   *
+   * @par Example
+   * @snippet samples.cc add-database-reader-on-database
+   *
+   * @param db the identifier for the database where you want to change the IAM
+   *     policy.
+   * @param updater a callback to modify the policy.  Return an unset optional
+   *     to indicate that no changes to the policy are needed.
+   */
+  StatusOr<google::iam::v1::Policy> SetIamPolicy(Database const& db,
+                                                 IamUpdater const& updater);
+
+  /**
+   * @copydoc SetIamPolicy(Database const&,IamUpdater const&)
+   *
+   * @param rerun_policy controls for how long (or how many times) the updater
+   *     will be rerun after the IAM policy update aborts.
+   * @param backoff_policy controls how long `SetIamPolicy` waits between
+   *     reruns.
+   */
+  StatusOr<google::iam::v1::Policy> SetIamPolicy(
+      Database const& db, IamUpdater const& updater,
+      std::unique_ptr<TransactionRerunPolicy> rerun_policy,
+      std::unique_ptr<BackoffPolicy> backoff_policy);
+
   /**
    * Get the subset of the permissions the caller has on the given database.
    *

google/cloud/spanner/database_admin_client_test.cc

@@ -25,7 +25,9 @@ namespace {
 
 using ::google::cloud::spanner_mocks::MockDatabaseAdminConnection;
 using ::testing::_;
+using ::testing::AtLeast;
 using ::testing::ElementsAre;
+using ::testing::HasSubstr;
 namespace gcsa = ::google::spanner::admin::database::v1;
 
 /// @test Verify DatabaseAdminClient uses CreateDatabase() correctly.
@@ -186,10 +188,130 @@ TEST(DatabaseAdminClientTest, SetIamPolicy) {
             return p.policy;
           });
   DatabaseAdminClient client(std::move(mock));
-  auto response = client.SetIamPolicy(expected_db, {});
+  auto response = client.SetIamPolicy(expected_db, google::iam::v1::Policy{});
   EXPECT_STATUS_OK(response);
 }
 
+TEST(DatabaseAdminClientTest, SetIamPolicyOccGetFailure) {
+  Database const db("test-project", "test-instance", "test-database");
+  auto mock = std::make_shared<MockDatabaseAdminConnection>();
+  EXPECT_CALL(*mock, GetIamPolicy(_))
+      .WillOnce([&db](DatabaseAdminConnection::GetIamPolicyParams const& p) {
+        EXPECT_EQ(db, p.database);
+        return Status(StatusCode::kPermissionDenied, "uh-oh");
+      });
+
+  DatabaseAdminClient client(mock);
+  auto actual = client.SetIamPolicy(db, [](google::iam::v1::Policy const&) {
+    return optional<google::iam::v1::Policy>{};
+  });
+  EXPECT_EQ(StatusCode::kPermissionDenied, actual.status().code());
+}
+
+TEST(DatabaseAdminClientTest, SetIamPolicyOccNoUpdates) {
+  Database const db("test-project", "test-instance", "test-database");
+  auto mock = std::make_shared<MockDatabaseAdminConnection>();
+  EXPECT_CALL(*mock, GetIamPolicy(_))
+      .WillOnce([&db](DatabaseAdminConnection::GetIamPolicyParams const& p) {
+        EXPECT_EQ(db, p.database);
+        google::iam::v1::Policy r;
+        r.set_etag("test-etag");
+        return r;
+      });
+  EXPECT_CALL(*mock, SetIamPolicy(_)).Times(0);
+
+  DatabaseAdminClient client(mock);
+  auto actual = client.SetIamPolicy(db, [](google::iam::v1::Policy const& p) {
+    EXPECT_EQ("test-etag", p.etag());
+    return optional<google::iam::v1::Policy>{};
+  });
+  ASSERT_STATUS_OK(actual);
+  EXPECT_EQ("test-etag", actual->etag());
+}
+
+std::unique_ptr<TransactionRerunPolicy> RerunPolicyForTesting() {
+  return LimitedErrorCountTransactionRerunPolicy(/*maximum_failures=*/3)
+      .clone();
+}
+
+std::unique_ptr<BackoffPolicy> BackoffPolicyForTesting() {
+  return ExponentialBackoffPolicy(
+             /*initial_delay=*/std::chrono::microseconds(1),
+             /*maximum_delay=*/std::chrono::microseconds(1), /*scaling=*/2.0)
+      .clone();
+}
+
+TEST(DatabaseAdminClientTest, SetIamPolicyOccRetryAborted) {
+  Database const db("test-project", "test-instance", "test-database");
+  auto mock = std::make_shared<MockDatabaseAdminConnection>();
+  EXPECT_CALL(*mock, GetIamPolicy(_))
+      .WillOnce([&db](DatabaseAdminConnection::GetIamPolicyParams const& p) {
+        EXPECT_EQ(db, p.database);
+        google::iam::v1::Policy r;
+        r.set_etag("test-etag-1");
+        return r;
+      })
+      .WillOnce([&db](DatabaseAdminConnection::GetIamPolicyParams const& p) {
+        EXPECT_EQ(db, p.database);
+        google::iam::v1::Policy r;
+        r.set_etag("test-etag-2");
+        return r;
+      });
+  EXPECT_CALL(*mock, SetIamPolicy(_))
+      .WillOnce([&db](DatabaseAdminConnection::SetIamPolicyParams const& p) {
+        EXPECT_EQ(db, p.database);
+        EXPECT_EQ("test-etag-1", p.policy.etag());
+        return Status(StatusCode::kAborted, "aborted");
+      })
+      .WillOnce([&db](DatabaseAdminConnection::SetIamPolicyParams const& p) {
+        EXPECT_EQ(db, p.database);
+        EXPECT_EQ("test-etag-2", p.policy.etag());
+        google::iam::v1::Policy r;
+        r.set_etag("test-etag-3");
+        return r;
+      });
+
+  DatabaseAdminClient client(mock);
+  int counter = 0;
+  auto actual = client.SetIamPolicy(
+      db,
+      [&counter](google::iam::v1::Policy p) {
+        EXPECT_EQ("test-etag-" + std::to_string(++counter), p.etag());
+        return p;
+      },
+      RerunPolicyForTesting(), BackoffPolicyForTesting());
+  ASSERT_STATUS_OK(actual);
+  EXPECT_EQ("test-etag-3", actual->etag());
+}
+
+TEST(DatabaseAdminClientTest, SetIamPolicyOccRetryAbortedTooManyFailures) {
+  Database const db("test-project", "test-instance", "test-database");
+  auto mock = std::make_shared<MockDatabaseAdminConnection>();
+  EXPECT_CALL(*mock, GetIamPolicy(_))
+      .WillRepeatedly(
+          [&db](DatabaseAdminConnection::GetIamPolicyParams const& p) {
+            EXPECT_EQ(db, p.database);
+            google::iam::v1::Policy r;
+            r.set_etag("test-etag-1");
+            return r;
+          });
+  EXPECT_CALL(*mock, SetIamPolicy(_))
+      .Times(AtLeast(2))
+      .WillRepeatedly(
+          [&db](DatabaseAdminConnection::SetIamPolicyParams const& p) {
+            EXPECT_EQ(db, p.database);
+            EXPECT_EQ("test-etag-1", p.policy.etag());
+            return Status(StatusCode::kAborted, "test-msg");
+          });
+
+  DatabaseAdminClient client(mock);
+  auto actual = client.SetIamPolicy(
+      db, [](google::iam::v1::Policy p) { return p; }, RerunPolicyForTesting(),
+      BackoffPolicyForTesting());
+  EXPECT_EQ(StatusCode::kAborted, actual.status().code());
+  EXPECT_THAT(actual.status().message(), HasSubstr("test-msg"));
+}
+
 /// @test Verify DatabaseAdminClient uses TestIamPermissions() correctly.
 TEST(DatabaseAdminClientTest, TestIamPermissions) {
   auto mock = std::make_shared<MockDatabaseAdminConnection>();

google/cloud/spanner/iam_updater.h

@@ -0,0 +1,35 @@
+// Copyright 2019 Google LLC
+//
+// Licensed 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.
+
+#ifndef GOOGLE_CLOUD_CPP_SPANNER_GOOGLE_CLOUD_SPANNER_IAM_UPDATER_H_
+#define GOOGLE_CLOUD_CPP_SPANNER_GOOGLE_CLOUD_SPANNER_IAM_UPDATER_H_
+
+#include "google/cloud/optional.h"
+#include <google/iam/v1/policy.pb.h>
+#include <functional>
+
+namespace google {
+namespace cloud {
+namespace spanner {
+inline namespace SPANNER_CLIENT_NS {
+
+using IamUpdater =
+    std::function<optional<google::iam::v1::Policy>(google::iam::v1::Policy)>;
+
+}  // namespace SPANNER_CLIENT_NS
+}  // namespace spanner
+}  // namespace cloud
+}  // namespace google
+
+#endif  // GOOGLE_CLOUD_CPP_SPANNER_GOOGLE_CLOUD_SPANNER_IAM_UPDATER_H_

google/cloud/spanner/instance_admin_client.cc

@@ -67,6 +67,54 @@ StatusOr<google::iam::v1::Policy> InstanceAdminClient::SetIamPolicy(
   return conn_->SetIamPolicy({in.FullName(), std::move(policy)});
 }
 
+StatusOr<google::iam::v1::Policy> InstanceAdminClient::SetIamPolicy(
+    Instance const& in, IamUpdater const& updater) {
+  auto const rerun_maximum_duration = std::chrono::minutes(15);
+  auto default_rerun_policy =
+      LimitedTimeTransactionRerunPolicy(rerun_maximum_duration).clone();
+
+  auto const backoff_initial_delay = std::chrono::milliseconds(1000);
+  auto const backoff_maximum_delay = std::chrono::minutes(5);
+  auto const backoff_scaling = 2.0;
+  auto default_backoff_policy =
+      ExponentialBackoffPolicy(backoff_initial_delay, backoff_maximum_delay,
+                               backoff_scaling)
+          .clone();
+
+  return SetIamPolicy(in, updater, std::move(default_rerun_policy),
+                      std::move(default_backoff_policy));
+}
+
+StatusOr<google::iam::v1::Policy> InstanceAdminClient::SetIamPolicy(
+    Instance const& in, IamUpdater const& updater,
+    std::unique_ptr<TransactionRerunPolicy> rerun_policy,
+    std::unique_ptr<BackoffPolicy> backoff_policy) {
+  using RerunnablePolicy = internal::SafeTransactionRerun;
+
+  Status last_status;
+  do {
+    auto current_policy = GetIamPolicy(in);
+    if (!current_policy) {
+      last_status = std::move(current_policy).status();
+    } else {
+      auto etag = current_policy->etag();
+      auto desired = updater(*current_policy);
+      if (!desired.has_value()) {
+        return current_policy;
+      }
+      desired->set_etag(std::move(etag));
+      auto result = SetIamPolicy(in, *std::move(desired));
+      if (RerunnablePolicy::IsOk(result.status())) {
+        return result;
+      }
+      last_status = std::move(result).status();
+    }
+    if (!rerun_policy->OnFailure(last_status)) break;
+    std::this_thread::sleep_for(backoff_policy->OnCompletion());
+  } while (!rerun_policy->IsExhausted());
+  return last_status;
+}
+
 StatusOr<google::iam::v1::TestIamPermissionsResponse>
 InstanceAdminClient::TestIamPermissions(Instance const& in,
                                         std::vector<std::string> permissions) {