From 16669025be42e67f7c0a4d383d48a6ed6e88d0e4 Mon Sep 17 00:00:00 2001
From: Watson <watson1978@gmail.com>
Date: Sun, 12 Nov 2023 16:14:37 +0900
Subject: [PATCH] Raise exception if await call twice

---
 .rubocop.yml        |  3 +++
 ext/ilios/future.c  |  5 +++++
 ext/ilios/ilios.h   |  1 +
 ext/ilios/session.c |  2 ++
 test/test_future.rb | 11 ++++++-----
 5 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/.rubocop.yml b/.rubocop.yml
index 531df40..52a611d 100644
--- a/.rubocop.yml
+++ b/.rubocop.yml
@@ -20,6 +20,9 @@ Layout/RedundantLineBreak:
 Lint/ConstantResolution:
   Enabled: false
 
+Lint/EmptyBlock:
+  Enabled: false
+
 Metrics/AbcSize:
   Enabled: false
 
diff --git a/ext/ilios/future.c b/ext/ilios/future.c
index 7c613e6..54d636d 100644
--- a/ext/ilios/future.c
+++ b/ext/ilios/future.c
@@ -254,11 +254,16 @@ static VALUE future_await(VALUE self)
 
     GET_FUTURE(self, cassandra_future);
 
+    if (cassandra_future->already_waited) {
+        rb_raise(eExecutionError, "It should not call twice");
+    }
+
     if (cassandra_future->on_success_block || cassandra_future->on_failure_block) {
         nogvl_sem_wait(&cassandra_future->sem);
     } else {
         nogvl_future_wait(cassandra_future->future);
     }
+    cassandra_future->already_waited = true;
     return self;
 }
 
diff --git a/ext/ilios/ilios.h b/ext/ilios/ilios.h
index 96ef7df..4b84646 100644
--- a/ext/ilios/ilios.h
+++ b/ext/ilios/ilios.h
@@ -56,6 +56,7 @@ typedef struct
     VALUE proc_mutex;
 
     uv_sem_t sem;
+    bool already_waited;
 } CassandraFuture;
 
 extern const rb_data_type_t cassandra_session_data_type;
diff --git a/ext/ilios/session.c b/ext/ilios/session.c
index e80fa96..fda1d69 100644
--- a/ext/ilios/session.c
+++ b/ext/ilios/session.c
@@ -39,6 +39,7 @@ static VALUE session_prepare_async(VALUE self, VALUE query)
     cassandra_future->session_obj = self;
     cassandra_future->proc_mutex = rb_mutex_new();
     uv_sem_init(&cassandra_future->sem, 0);
+    cassandra_future->already_waited = false;
 
     return cassandra_future_obj;
 }
@@ -109,6 +110,7 @@ static VALUE session_execute_async(VALUE self, VALUE statement)
     cassandra_future->statement_obj = statement;
     cassandra_future->proc_mutex = rb_mutex_new();
     uv_sem_init(&cassandra_future->sem, 0);
+    cassandra_future->already_waited = false;
 
     return cassandra_future_obj;
 }
diff --git a/test/test_future.rb b/test/test_future.rb
index 544439a..7ebf074 100644
--- a/test/test_future.rb
+++ b/test/test_future.rb
@@ -52,6 +52,7 @@ def test_await
     prepare_future.await
 
     assert_equal(50, count)
+    assert_raises(Ilios::Cassandra::ExecutionError) { prepare_future.await }
   end
 
   def test_on_success
@@ -60,9 +61,9 @@ def test_on_success
 
     assert_raises(ArgumentError) { future.on_success }
 
-    future.on_success { }
+    future.on_success {}
 
-    assert_raises(Ilios::Cassandra::ExecutionError) { future.on_success { } }
+    assert_raises(Ilios::Cassandra::ExecutionError) { future.on_success {} }
   end
 
   def test_on_failure
@@ -71,8 +72,8 @@ def test_on_failure
 
     assert_raises(ArgumentError) { future.on_failure }
 
-    future.on_failure { }
+    future.on_failure {}
 
-    assert_raises(Ilios::Cassandra::ExecutionError) { future.on_failure { } }
-end
+    assert_raises(Ilios::Cassandra::ExecutionError) { future.on_failure {} }
+  end
 end