Restore reader_count in case of failure in rwlock_rdlock_impl

Author: heshan

src/bthread/rwlock.cpp

@@ -41,7 +41,8 @@ const int RWLockMaxReaders = 1 << 30;
 // For reading.
 static int rwlock_rdlock_impl(bthread_rwlock_t* __restrict rwlock,
                               const struct timespec* __restrict abstime) {
-    int reader_count = ((butil::atomic<int>*)&rwlock->reader_count)
+    auto* reader_count_atomic = (butil::atomic<int>*)&rwlock->reader_count;
+    int reader_count = reader_count_atomic
         ->fetch_add(1, butil::memory_order_acquire) + 1;
     // Fast path.
     if (reader_count >= 0) {
@@ -55,7 +56,7 @@ static int rwlock_rdlock_impl(bthread_rwlock_t* __restrict rwlock,
     if (NULL == bthread::g_cp) {
         int rc = bthread_sem_timedwait(&rwlock->reader_sema, abstime);
         if (rc != 0){
-            ((butil::atomic<int>*)&rwlock->reader_count)->fetch_add(-1, butil::memory_order_relaxed);
+            reader_count_atomic->fetch_add(-1, butil::memory_order_relaxed);
         }
         return rc;
     }
@@ -64,14 +65,17 @@ static int rwlock_rdlock_impl(bthread_rwlock_t* __restrict rwlock,
     if (!bvar::is_sampling_range_valid(sampling_range)) { // Don't sample.
         int rc = bthread_sem_timedwait(&rwlock->reader_sema, abstime);
         if (rc != 0){
-            ((butil::atomic<int>*)&rwlock->reader_count)->fetch_add(-1, butil::memory_order_relaxed);
+            reader_count_atomic->fetch_add(-1, butil::memory_order_relaxed);
         }
         return rc;
     }
 
     // Sample.
     const int64_t start_ns = butil::cpuwide_time_ns();
     int rc = bthread_sem_timedwait(&rwlock->reader_sema, abstime);
+    if (rc != 0){
+        reader_count_atomic->fetch_add(-1, butil::memory_order_relaxed);
+    }
     const int64_t end_ns = butil::cpuwide_time_ns();
     if (rc != 0){
         ((butil::atomic<int>*)&rwlock->reader_count)->fetch_add(-1, butil::memory_order_relaxed);
@@ -200,7 +204,8 @@ static inline int rwlock_wrlock_impl(bthread_rwlock_t* __restrict rwlock,
     }
 
     // Announce to readers there is a pending writer.
-    int reader_count = ((butil::atomic<int>*)&rwlock->reader_count)
+    auto* reader_count_atomic = (butil::atomic<int>*)&rwlock->reader_count;
+    int reader_count = reader_count_atomic
         ->fetch_add(-RWLockMaxReaders, butil::memory_order_release);
     // Wait for active readers.
     if (reader_count != 0 &&
@@ -215,6 +220,7 @@ static inline int rwlock_wrlock_impl(bthread_rwlock_t* __restrict rwlock,
             rc = bthread_sem_timedwait(&rwlock->writer_sema, abstime);
             if (0 != rc) {
                 SUBMIT_CSITE_IF_NEED;
+                reader_count_atomic->fetch_add(RWLockMaxReaders, butil::memory_order_release);
                 bthread_mutex_unlock(&rwlock->write_queue_mutex);
                 return rc;
             }

test/BUILD.bazel

@@ -238,7 +238,7 @@ cc_test(
             # glog CHECK die with a fatal error
             "bthread_key_unittest.cpp",
             "bthread_butex_multi_tag_unittest.cpp",
-            "bthread_rwlock_unittest.cpp",
+            #"bthread_rwlock_unittest.cpp",
             "bthread_semaphore_unittest.cpp",
         ],
     ),

test/bthread_rwlock_unittest.cpp

@@ -85,24 +85,23 @@ TEST(RWLockTest, used_in_pthread) {
 }
 
 void* do_timedrdlock(void *arg) {
-    struct timespec t = { -2, 0 };
+    struct timespec t = { 0, 100 };
     EXPECT_EQ(ETIMEDOUT, bthread_rwlock_timedrdlock((bthread_rwlock_t*)arg, &t));
     return NULL;
 }
 
 void* do_timedwrlock(void *arg) {
-    struct timespec t = { -2, 0 };
+    struct timespec t = { 0, 100 };
     EXPECT_EQ(ETIMEDOUT, bthread_rwlock_timedwrlock((bthread_rwlock_t*)arg, &t));
-    LOG(INFO) << 10;
     return NULL;
 }
 
 TEST(RWLockTest, timedlock) {
     bthread_rwlock_t rw;
     ASSERT_EQ(0, bthread_rwlock_init(&rw, NULL));
+    bthread_t th;
 
     ASSERT_EQ(0, bthread_rwlock_rdlock(&rw));
-    bthread_t th;
     ASSERT_EQ(0, bthread_start_urgent(&th, NULL, do_timedwrlock, &rw));
     ASSERT_EQ(0, bthread_join(th, NULL));
     ASSERT_EQ(0, bthread_rwlock_unlock(&rw));
@@ -113,6 +112,19 @@ TEST(RWLockTest, timedlock) {
     ASSERT_EQ(0, bthread_start_urgent(&th, NULL, do_timedrdlock, &rw));
     ASSERT_EQ(0, bthread_join(th, NULL));
     ASSERT_EQ(0, bthread_rwlock_unlock(&rw));
+
+    ASSERT_EQ(0, bthread_rwlock_wrlock(&rw));
+    ASSERT_EQ(0, bthread_start_urgent(&th, NULL, do_timedrdlock, &rw));
+    ASSERT_EQ(0, bthread_join(th, NULL));
+    ASSERT_EQ(0, bthread_start_urgent(&th, NULL, do_timedwrlock, &rw));
+    ASSERT_EQ(0, bthread_join(th, NULL));
+    ASSERT_EQ(0, bthread_rwlock_unlock(&rw));
+
+    ASSERT_EQ(0, bthread_rwlock_rdlock(&rw));
+    ASSERT_EQ(0, bthread_start_urgent(&th, NULL, do_timedwrlock, &rw));
+    ASSERT_EQ(0, bthread_join(th, NULL));
+    ASSERT_EQ(0, bthread_rwlock_unlock(&rw));
+
     ASSERT_EQ(0, bthread_rwlock_destroy(&rw));
 }