diff --git a/lib/pls/src/internal/base/tas_spin_lock.cpp b/lib/pls/src/internal/base/tas_spin_lock.cpp
index b35d805..5132b25 100644
--- a/lib/pls/src/internal/base/tas_spin_lock.cpp
+++ b/lib/pls/src/internal/base/tas_spin_lock.cpp
@@ -8,26 +8,34 @@ namespace base {
 void tas_spin_lock::lock() {
   PROFILE_LOCK("Acquire Lock")
   int tries = 0;
-  while (flag_.test_and_set(std::memory_order_acquire)) {
+  while (true) {
     tries++;
     if (tries % yield_at_tries_ == 0) {
       this_thread::yield();
     }
+
+    if (flag_.test_and_set(std::memory_order_acquire) == 0) {
+      return;
+    }
   }
 }
 
 bool tas_spin_lock::try_lock(unsigned int num_tries) {
   PROFILE_LOCK("Try Acquire Lock")
-  while (flag_.test_and_set(std::memory_order_acquire)) {
+  while (true) {
     num_tries--;
     if (num_tries <= 0) {
       return false;
     }
+
+    if (flag_.test_and_set(std::memory_order_acquire) == 0) {
+      return true;
+    }
   }
-  return true;
 }
 
 void tas_spin_lock::unlock() {
+  PROFILE_LOCK("Unlock")
   flag_.clear(std::memory_order_release);
 }
 
diff --git a/lib/pls/src/internal/base/ttas_spin_lock.cpp b/lib/pls/src/internal/base/ttas_spin_lock.cpp
index ba1d3ac..806251a 100644
--- a/lib/pls/src/internal/base/ttas_spin_lock.cpp
+++ b/lib/pls/src/internal/base/ttas_spin_lock.cpp
@@ -10,7 +10,7 @@ void ttas_spin_lock::lock() {
   int tries = 0;
   int expected = 0;
 
-  do {
+  while (true) {
     while (flag_.load(std::memory_order_relaxed) == 1) {
       tries++;
       if (tries % yield_at_tries_ == 0) {
@@ -19,14 +19,17 @@ void ttas_spin_lock::lock() {
     }
 
     expected = 0;
-  } while (!flag_.compare_exchange_weak(expected, 1, std::memory_order_acquire));
+    if (flag_.compare_exchange_weak(expected, 1, std::memory_order_acquire)) {
+      return;
+    }
+  }
 }
 
 bool ttas_spin_lock::try_lock(unsigned int num_tries) {
   PROFILE_LOCK("Try Acquire Lock")
   int expected = 0;
 
-  do {
+  while (true) {
     while (flag_.load(std::memory_order_relaxed) == 1) {
       num_tries--;
       if (num_tries <= 0) {
@@ -35,12 +38,14 @@ bool ttas_spin_lock::try_lock(unsigned int num_tries) {
     }
 
     expected = 0;
-  } while (!flag_.compare_exchange_weak(expected, 1, std::memory_order_acquire));
-
-  return true;
+    if (flag_.compare_exchange_weak(expected, 1, std::memory_order_acquire)) {
+      return true;
+    }
+  }
 }
 
 void ttas_spin_lock::unlock() {
+  PROFILE_LOCK("Unlock")
   flag_.store(0, std::memory_order_release);
 }