diff --git a/src/client_context.cpp b/src/client_context.cpp
index 47200d6..9dfc9c0 100644
--- a/src/client_context.cpp
+++ b/src/client_context.cpp
@@ -12,9 +12,10 @@ trrep::provider& trrep::client_context::provider() const
 
 // TODO: This should be pure virtual method, implemented by
 // DBMS integration or mock classes only.
-int trrep::client_context::replay(trrep::transaction_context& tc)
+int trrep::client_context::replay(
+    trrep::unique_lock<trrep::mutex>& lock,
+    trrep::transaction_context& tc)
 {
-    trrep::unique_lock<trrep::mutex> lock(mutex_);
     tc.state(lock, trrep::transaction_context::s_replaying);
     tc.state(lock, trrep::transaction_context::s_committed);
     return 0;
diff --git a/src/client_context.hpp b/src/client_context.hpp
index d5f9873..4bd3b0c 100644
--- a/src/client_context.hpp
+++ b/src/client_context.hpp
@@ -84,7 +84,8 @@ namespace trrep
         virtual int rollback(trrep::transaction_context&) = 0;
 
         virtual void will_replay(trrep::transaction_context&) { }
-        virtual int replay(trrep::transaction_context& tc);
+        virtual int replay(trrep::unique_lock<trrep::mutex>&,
+                           trrep::transaction_context& tc);
 
 
         virtual int apply(trrep::transaction_context&,
diff --git a/src/transaction_context.cpp b/src/transaction_context.cpp
index 8a26388..099301e 100644
--- a/src/transaction_context.cpp
+++ b/src/transaction_context.cpp
@@ -70,6 +70,7 @@ int trrep::transaction_context::start_transaction(
 {
     assert(active() == false);
     id_ = id;
+    state_ = s_executing;
     ws_handle_.trx_id = id_.get();
     flags_ |= WSREP_FLAG_TRX_START;
     switch (client_context_.mode())
@@ -419,7 +420,7 @@ int trrep::transaction_context::after_statement()
     case s_aborted:
         break;
     case s_must_replay:
-        ret = client_context_.replay(*this);
+        ret = client_context_.replay(lock, *this);
         break;
     default:
         assert(0);
@@ -703,7 +704,10 @@ void trrep::transaction_context::cleanup()
     //        << client_context_.id().get()
     //         << ": " << id_.get() << "\n";
     id_ = trrep::transaction_id::invalid();
-    state_ = s_executing;
+    if (is_streaming())
+    {
+        state_ = s_executing;
+    }
     state_hist_.clear();
     trx_meta_.gtid = WSREP_GTID_UNDEFINED;
     trx_meta_.stid.node = WSREP_UUID_UNDEFINED;