From 7ec0989eea7e56990cdf4b6f7b4ce14ca21ec91f Mon Sep 17 00:00:00 2001
From: "Dunajski, Bartosz" <bartosz.dunajski@intel.com>
Date: Mon, 10 Sep 2018 14:24:11 +0200
Subject: [PATCH] Obtain new TimestampPacket before releasing old node

Change-Id: I243d50e32967a0accea2d73fdd5ff4935bc03449
Signed-off-by: Dunajski, Bartosz <bartosz.dunajski@intel.com>
---
 runtime/command_queue/command_queue.cpp       |  7 +++---
 unit_tests/helpers/timestamp_packet_tests.cpp | 23 +++++++++++++++++++
 unit_tests/mocks/mock_command_queue.h         |  2 ++
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/runtime/command_queue/command_queue.cpp b/runtime/command_queue/command_queue.cpp
index b36bc53b56..2012e7d368 100644
--- a/runtime/command_queue/command_queue.cpp
+++ b/runtime/command_queue/command_queue.cpp
@@ -599,9 +599,10 @@ void CommandQueue::dispatchAuxTranslation(MultiDispatchInfo &multiDispatchInfo,
 void CommandQueue::obtainNewTimestampPacketNode() {
     auto allocator = device->getMemoryManager()->getTimestampPacketAllocator();
 
-    if (timestampPacketNode) {
-        allocator->returnTag(timestampPacketNode);
-    }
+    auto oldNode = timestampPacketNode;
     timestampPacketNode = allocator->getTag();
+    if (oldNode) {
+        allocator->returnTag(oldNode);
+    }
 }
 } // namespace OCLRT
diff --git a/unit_tests/helpers/timestamp_packet_tests.cpp b/unit_tests/helpers/timestamp_packet_tests.cpp
index 0e8d2eaca9..fcfe16a4e9 100644
--- a/unit_tests/helpers/timestamp_packet_tests.cpp
+++ b/unit_tests/helpers/timestamp_packet_tests.cpp
@@ -46,6 +46,7 @@ struct TimestampPacketTests : public ::testing::Test {
     class MockTagAllocator : public TagAllocator<TagType> {
       public:
         using BaseClass = TagAllocator<TagType>;
+        using BaseClass::freeTags;
         using BaseClass::usedTags;
         using NodeType = typename BaseClass::NodeType;
 
@@ -479,3 +480,25 @@ HWTEST_F(TimestampPacketTests, givenTimestampPacketWriteEnabledWhenDispatchingTh
     EXPECT_EQ(1u, walkersFound);
     EXPECT_EQ(2u, semaphoresFound); // total number of semaphores found in cmdList
 }
+
+TEST_F(TimestampPacketTests, givenAlreadyAssignedNodeWhenObtainingThenGetNewBeforeReleasing) {
+    auto device = std::unique_ptr<MockDevice>(MockDevice::createWithNewExecutionEnvironment<MockDevice>(platformDevices[0]));
+    auto mockMemoryManager = new MockMemoryManager();
+    device->injectMemoryManager(mockMemoryManager);
+    auto mockTagAllocator = new MockTagAllocator<>(mockMemoryManager, 1);
+    mockMemoryManager->timestampPacketAllocator.reset(mockTagAllocator);
+
+    MockCommandQueue cmdQ(nullptr, device.get(), nullptr);
+    cmdQ.obtainNewTimestampPacketNode();
+    auto firstNode = cmdQ.timestampPacketNode;
+    EXPECT_TRUE(mockTagAllocator->freeTags.peekIsEmpty());
+
+    // mark as ready to release
+    size_t dataSize = sizeof(uint32_t) * static_cast<size_t>(TimestampPacket::DataIndex::Max);
+    memset(reinterpret_cast<void *>(firstNode->tag->pickAddressForDataWrite(TimestampPacket::DataIndex::ContextStart)), 0, dataSize);
+
+    cmdQ.obtainNewTimestampPacketNode();
+    auto secondNode = cmdQ.timestampPacketNode;
+    EXPECT_FALSE(mockTagAllocator->freeTags.peekIsEmpty()); // new pool allocated for secondNode
+    EXPECT_NE(firstNode, secondNode);
+}
diff --git a/unit_tests/mocks/mock_command_queue.h b/unit_tests/mocks/mock_command_queue.h
index 36dffe815d..0e1bdb66c5 100644
--- a/unit_tests/mocks/mock_command_queue.h
+++ b/unit_tests/mocks/mock_command_queue.h
@@ -31,6 +31,8 @@ namespace OCLRT {
 class MockCommandQueue : public CommandQueue {
   public:
     using CommandQueue::device;
+    using CommandQueue::obtainNewTimestampPacketNode;
+    using CommandQueue::timestampPacketNode;
 
     void setProfilingEnabled() {
         commandQueueProperties |= CL_QUEUE_PROFILING_ENABLE;