Fix timer_handler memory leak

Author: Jinming-Hu

Microsoft.WindowsAzure.Storage/includes/wascore/timer_handler.h

@@ -45,38 +45,43 @@ namespace azure { namespace storage { namespace core {
 
         WASTORAGE_API void stop_timer();
 
-        pplx::cancellation_token get_cancellation_token()
+        bool timer_started() const
         {
-            return m_worker_cancellation_token_source->get_token();
+            return m_timer_started.load(std::memory_order_acquire);
         }
 
-        bool is_canceled()
+        pplx::cancellation_token get_cancellation_token() const
         {
-            return m_worker_cancellation_token_source->get_token().is_canceled();
+            return m_worker_cancellation_token_source.get_token();
         }
 
-        bool is_canceled_by_timeout()
+        bool is_canceled() const
         {
-            return m_is_canceled_by_timeout;
+            return m_worker_cancellation_token_source.get_token().is_canceled();
+        }
+
+        bool is_canceled_by_timeout() const
+        {
+            return m_is_canceled_by_timeout.load(std::memory_order_acquire);
         }
 
     private:
-        std::shared_ptr<pplx::cancellation_token_source> m_worker_cancellation_token_source;
+        pplx::cancellation_token_source m_worker_cancellation_token_source;
         pplx::cancellation_token_registration m_cancellation_token_registration;
         pplx::cancellation_token m_cancellation_token;
         pplx::task<void> m_timeout_task;
-        bool m_is_canceled_by_timeout;
+        std::atomic<bool> m_is_canceled_by_timeout;
         pplx::task_completion_event<void> m_tce;
 
-        std::shared_ptr<std::mutex> m_mutex;
+        std::mutex m_mutex;
 
         WASTORAGE_API pplx::task<void> timeout_after(const std::chrono::milliseconds& time);
 
 #ifndef _WIN32
-        typedef std::chrono::steady_clock std_clock;
-        std::shared_ptr<boost::asio::basic_waitable_timer<std_clock>> m_timer;
+        std::shared_ptr<boost::asio::basic_waitable_timer<std::chrono::steady_clock>> m_timer;
 #else
         std::shared_ptr<concurrency::timer<int>> m_timer;
 #endif
+        std::atomic<bool> m_timer_started;
     };
 }}}

Microsoft.WindowsAzure.Storage/src/executor.cpp

@@ -32,8 +32,8 @@ namespace azure { namespace storage { namespace core {
         auto instance = std::make_shared<executor_impl>(command, options, context);
         return pplx::details::_do_while([instance]() -> pplx::task<bool>
         {
-            //Start the timer to track timeout.
-            if (instance->m_command->m_use_timeout)
+            // Start the timer to track timeout.
+            if (instance->m_command->m_use_timeout && !instance->m_command->m_timer_handler->timer_started())
             {
                 // Timer will be stopped when instance is out of scope, so no need to stop here.
                 instance->m_command->m_timer_handler->start_timer(instance->m_request_options.maximum_execution_time());

Microsoft.WindowsAzure.Storage/src/timer_handler.cpp

@@ -21,15 +21,14 @@
 namespace azure {    namespace storage {    namespace core {
 
     timer_handler::timer_handler(const pplx::cancellation_token& token) :
-        m_cancellation_token(token), m_is_canceled_by_timeout(false)
+        m_cancellation_token(token), m_is_canceled_by_timeout(false), m_timer_started(false)
     {
-        m_worker_cancellation_token_source = std::make_shared<pplx::cancellation_token_source>();
         if (m_cancellation_token != pplx::cancellation_token::none())
         {
             m_cancellation_token_registration = m_cancellation_token.register_callback([this]()
             {
-                this->m_worker_cancellation_token_source->cancel();
-                this->stop_timer();
+                m_worker_cancellation_token_source.cancel();
+                stop_timer();
             });
         }
     }
@@ -41,67 +40,66 @@ namespace azure {    namespace storage {    namespace core {
             m_cancellation_token.deregister_callback(m_cancellation_token_registration);
         }
 
-#ifdef _WIN32
         stop_timer();
-#else // LINUX
-        try
-        {
-            stop_timer();
-        }
-        catch (boost::exception_detail::clone_impl<boost::exception_detail::error_info_injector<boost::system::system_error> > &e)
-        {
-        }
-#endif
-
     }
 
     void timer_handler::start_timer(const std::chrono::milliseconds& time)
     {
-        m_mutex = std::make_shared<std::mutex>();
-        auto this_pointer = std::dynamic_pointer_cast<timer_handler>(shared_from_this());
-        m_timeout_task = timeout_after(time).then([this_pointer]()
+        std::lock_guard<std::mutex> guard(m_mutex);
+        if (m_timer_started.load(std::memory_order_acquire))
+        {
+            return;
+        }
+        m_timer_started.store(true, std::memory_order_release);
+        std::weak_ptr<timer_handler> weak_this_pointer = shared_from_this();
+        m_timeout_task = timeout_after(time).then([weak_this_pointer]()
         {
-            this_pointer->m_is_canceled_by_timeout = true;
-            this_pointer->m_worker_cancellation_token_source->cancel();
+            auto this_pointer = weak_this_pointer.lock();
+            if (this_pointer)
+            {
+                this_pointer->m_is_canceled_by_timeout.store(true, std::memory_order_release);
+                this_pointer->m_worker_cancellation_token_source.cancel();
+            }
         });
     }
 
     void timer_handler::stop_timer()
     {
-        if (m_timer != nullptr)
+        std::lock_guard<std::mutex> guard(m_mutex);
+        if (m_timer_started.load(std::memory_order_acquire) && m_timer)
         {
-            std::lock_guard<std::mutex> guard(*m_mutex);
-            if (m_timer != nullptr)
-            {
 #ifndef _WIN32
-                m_timer->cancel();
+            m_timer->cancel();
 #else
-                m_timer->stop();
+            m_timer->stop();
 #endif
-                if (!m_tce._IsTriggered())
-                {
-                    // if task_completion_event is not yet triggered, it means timeout has not been triggered.
-                    m_tce._Cancel();
-                }
-                m_timer.reset();
+            if (!m_tce._IsTriggered())
+            {
+                // If task_completion_event is not yet triggered, it means timeout has not been triggered.
+                m_tce._Cancel();
             }
+            m_timer.reset();
         }
     }
 
 #ifndef _WIN32
     pplx::task<void> timer_handler::timeout_after(const std::chrono::milliseconds& time)
     {
-        m_timer = std::make_shared<boost::asio::basic_waitable_timer<std_clock>>(crossplat::threadpool::shared_instance().service());
-        m_timer->expires_from_now(std::chrono::duration_cast<std_clock::duration>(time));
-        auto this_pointer = std::dynamic_pointer_cast<timer_handler>(shared_from_this());
-        auto callback = [this_pointer](const boost::system::error_code& ec)
+        m_timer = std::make_shared<boost::asio::basic_waitable_timer<std::chrono::steady_clock>>(crossplat::threadpool::shared_instance().service());
+        m_timer->expires_from_now(std::chrono::duration_cast<std::chrono::steady_clock::duration>(time));
+        std::weak_ptr<timer_handler> weak_this_pointer = shared_from_this();
+        auto callback = [weak_this_pointer](const boost::system::error_code& ec)
         {
             if (ec != boost::asio::error::operation_aborted)
             {
-                std::lock_guard<std::mutex> guard(*(this_pointer->m_mutex));
-                if (!this_pointer->m_tce._IsTriggered())
+                auto this_pointer = weak_this_pointer.lock();
+                if (this_pointer)
                 {
-                    this_pointer->m_tce.set();
+                    std::lock_guard<std::mutex> guard(this_pointer->m_mutex);
+                    if (!this_pointer->m_tce._IsTriggered())
+                    {
+                        this_pointer->m_tce.set();
+                    }
                 }
             }
         };
@@ -114,23 +112,27 @@ namespace azure {    namespace storage {    namespace core {
 #else
     pplx::task<void> timer_handler::timeout_after(const std::chrono::milliseconds& time)
     {
-        // initialize the timer and connect the callback with completion event.
+        // Initialize the timer and connect the callback with completion event.
         m_timer = std::make_shared<concurrency::timer<int>>(static_cast<unsigned int>(time.count()), 0);
-        auto this_pointer = std::dynamic_pointer_cast<timer_handler>(shared_from_this());
-        auto callback = std::make_shared<concurrency::call<int>>([this_pointer](int)
+        std::weak_ptr<timer_handler> weak_this_pointer = shared_from_this();
+        auto callback = std::make_shared<concurrency::call<int>>([weak_this_pointer](int)
         {
-            std::lock_guard<std::mutex> guard(*(this_pointer->m_mutex));
-            if (!this_pointer->m_tce._IsTriggered())
+            auto this_pointer = weak_this_pointer.lock();
+            if (this_pointer)
             {
-                this_pointer->m_tce.set();
+                std::lock_guard<std::mutex> guard(this_pointer->m_mutex);
+                if (!this_pointer->m_tce._IsTriggered())
+                {
+                    this_pointer->m_tce.set();
+                }
             }
         });
-        m_timer->link_target(callback.get());//When timer stops, tce will trigger cancellation.
+        m_timer->link_target(callback.get());  // When timer stops, tce will trigger cancellation.
         m_timer->start();
 
         auto event_set = pplx::create_task(m_tce);
 
-        //timer and callback should be preserved before event set has been triggered.
+        // Timer and callback should be preserved before event set has been triggered.
         return event_set.then([callback]() {});
     }
 #endif