Add missing locks to functions operating on AUB file stream

Related-To: NEO-2747

Change-Id: I9efacbaf6d7894943f3abb6ebe2634ac34fc3d04
Signed-off-by: Milczarek, Slawomir <slawomir.milczarek@intel.com>

runtime/aub_mem_dump/aub_mem_dump.h

@@ -139,11 +139,11 @@ struct AubFileStream : public AubStream {
     MOCKABLE_VIRTUAL void expectMemory(uint64_t physAddress, const void *memory, size_t size,
                                        uint32_t addressSpace, uint32_t compareOperation);
     MOCKABLE_VIRTUAL bool addComment(const char *message);
-    MOCKABLE_VIRTUAL std::unique_lock<std::mutex> lockStream();
+    MOCKABLE_VIRTUAL std::unique_lock<std::recursive_mutex> lockStream();
 
     std::ofstream fileHandle;
     std::string fileName;
-    std::mutex mutex;
+    std::recursive_mutex mutex;
 };
 
 template <int addressingBits>

runtime/command_stream/aub_command_stream_receiver.cpp

@@ -277,8 +277,8 @@ bool AubFileStream::addComment(const char *message) {
     return true;
 }
 
-std::unique_lock<std::mutex> AubFileStream::lockStream() {
-    return std::unique_lock<std::mutex>(mutex);
+std::unique_lock<std::recursive_mutex> AubFileStream::lockStream() {
+    return std::unique_lock<std::recursive_mutex>(mutex);
 }
 
 } // namespace AubMemDump

runtime/command_stream/aub_command_stream_receiver_hw.h

@@ -66,7 +66,7 @@ class AUBCommandStreamReceiverHw : public CommandStreamReceiverSimulatedHw<GfxFa
 
     uint32_t getDumpHandle();
     MOCKABLE_VIRTUAL void addContextToken(uint32_t dumpHandle);
-    void dumpAllocation(GraphicsAllocation &gfxAllocation);
+    MOCKABLE_VIRTUAL void dumpAllocation(GraphicsAllocation &gfxAllocation);
 
     static CommandStreamReceiver *create(const std::string &fileName, bool standalone, ExecutionEnvironment &executionEnvironment);
 

runtime/command_stream/aub_command_stream_receiver_hw.inl

@@ -554,6 +554,8 @@ void AUBCommandStreamReceiverHw<GfxFamily>::pollForCompletionImpl() {
         }
     }
 
+    auto streamLocked = getAubStream()->lockStream();
+
     if (hardwareContextController) {
         hardwareContextController->pollForCompletion();
         return;
@@ -663,6 +665,8 @@ cl_int AUBCommandStreamReceiverHw<GfxFamily>::expectMemory(const void *gfxAddres
                                                            size_t length, uint32_t compareOperation) {
     pollForCompletion();
 
+    auto streamLocked = getAubStream()->lockStream();
+
     if (hardwareContextController) {
         hardwareContextController->expectMemory(reinterpret_cast<uint64_t>(gfxAddress), srcAddress, length, compareOperation);
     }
@@ -723,6 +727,8 @@ void AUBCommandStreamReceiverHw<GfxFamily>::dumpAllocation(GraphicsAllocation &g
         pollForCompletion();
     }
 
+    auto streamLocked = getAubStream()->lockStream();
+
     if (hardwareContextController) {
         auto surfaceInfo = std::unique_ptr<aub_stream::SurfaceInfo>(AubAllocDump::getDumpSurfaceInfo<GfxFamily>(gfxAllocation, dumpFormat));
         if (nullptr != surfaceInfo) {
@@ -765,6 +771,7 @@ void AUBCommandStreamReceiverHw<GfxFamily>::activateAubSubCapture(const MultiDis
 
 template <typename GfxFamily>
 void AUBCommandStreamReceiverHw<GfxFamily>::addAubComment(const char *message) {
+    auto streamLocked = getAubStream()->lockStream();
     if (aubManager) {
         aubManager->addComment(message);
         return;

unit_tests/command_stream/aub_file_stream_tests.cpp

@@ -186,6 +186,54 @@ HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenFlushIsCalledThenF
     EXPECT_TRUE(mockAubFileStream->lockStreamCalled);
 }
 
+HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenPollForCompletionIsCalledThenFileStreamShouldBeLocked) {
+    auto mockAubFileStream = std::make_unique<MockAubFileStream>();
+    auto aubExecutionEnvironment = getEnvironment<MockAubCsr<FamilyType>>(true, true, true);
+    auto aubCsr = aubExecutionEnvironment->template getCsr<MockAubCsr<FamilyType>>();
+
+    aubCsr->stream = static_cast<MockAubFileStream *>(mockAubFileStream.get());
+
+    aubCsr->latestSentTaskCount = 1;
+    aubCsr->pollForCompletionTaskCount = 0;
+
+    aubCsr->pollForCompletion();
+    EXPECT_TRUE(mockAubFileStream->lockStreamCalled);
+}
+
+HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenExpectMemoryEqualIsCalledThenFileStreamShouldBeLocked) {
+    auto mockAubFileStream = std::make_unique<MockAubFileStream>();
+    auto aubExecutionEnvironment = getEnvironment<MockAubCsr<FamilyType>>(true, true, true);
+    auto aubCsr = aubExecutionEnvironment->template getCsr<MockAubCsr<FamilyType>>();
+
+    aubCsr->stream = static_cast<MockAubFileStream *>(mockAubFileStream.get());
+
+    aubCsr->expectMemoryEqual(reinterpret_cast<void *>(0x1000), reinterpret_cast<void *>(0x1000), 0x1000);
+    EXPECT_TRUE(mockAubFileStream->lockStreamCalled);
+}
+
+HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenAddAubCommentIsCalledThenFileStreamShouldBeLocked) {
+    auto mockAubFileStream = std::make_unique<MockAubFileStream>();
+    auto aubExecutionEnvironment = getEnvironment<MockAubCsr<FamilyType>>(true, true, true);
+    auto aubCsr = aubExecutionEnvironment->template getCsr<MockAubCsr<FamilyType>>();
+
+    aubCsr->stream = static_cast<MockAubFileStream *>(mockAubFileStream.get());
+
+    aubCsr->addAubComment("comment");
+    EXPECT_TRUE(mockAubFileStream->lockStreamCalled);
+}
+
+HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenDumpAllocationIsCalledThenFileStreamShouldBeLocked) {
+    auto mockAubFileStream = std::make_unique<MockAubFileStream>();
+    auto aubExecutionEnvironment = getEnvironment<MockAubCsr<FamilyType>>(true, true, true);
+    auto aubCsr = aubExecutionEnvironment->template getCsr<MockAubCsr<FamilyType>>();
+    GraphicsAllocation allocation{GraphicsAllocation::AllocationType::UNKNOWN, nullptr, 0, 0, 0, MemoryPool::MemoryNull, false};
+
+    aubCsr->stream = static_cast<MockAubFileStream *>(mockAubFileStream.get());
+
+    aubCsr->dumpAllocation(allocation);
+    EXPECT_TRUE(mockAubFileStream->lockStreamCalled);
+}
+
 HWTEST_F(AubFileStreamTests, givenAubCommandStreamReceiverWhenFlushIsCalledThenItShouldCallTheExpectedFunctions) {
     auto aubExecutionEnvironment = getEnvironment<MockAubCsr<FamilyType>>(true, true, true);
     auto aubCsr = aubExecutionEnvironment->template getCsr<MockAubCsr<FamilyType>>();

unit_tests/mocks/mock_aub_csr.h

@@ -118,6 +118,10 @@ struct MockAubCsr : public AUBCommandStreamReceiverHw<GfxFamily> {
         AUBCommandStreamReceiverHw<GfxFamily>::addAubComment(message);
         addAubCommentCalled = true;
     }
+    void dumpAllocation(GraphicsAllocation &gfxAllocation) override {
+        AUBCommandStreamReceiverHw<GfxFamily>::dumpAllocation(gfxAllocation);
+        dumpAllocationCalled = true;
+    }
     bool flushBatchedSubmissionsCalled = false;
     bool initProgrammingFlagsCalled = false;
     bool initializeEngineCalled = false;
@@ -128,6 +132,7 @@ struct MockAubCsr : public AUBCommandStreamReceiverHw<GfxFamily> {
     bool expectMemoryEqualCalled = false;
     bool expectMemoryNotEqualCalled = false;
     bool addAubCommentCalled = false;
+    bool dumpAllocationCalled = false;
 
     void initFile(const std::string &fileName) override {
         fileIsOpen = true;

unit_tests/mocks/mock_aub_file_stream.h

@@ -42,7 +42,7 @@ struct MockAubFileStream : public AUBCommandStreamReceiver::AubFileStream {
     void flush() override {
         flushCalled = true;
     }
-    std::unique_lock<std::mutex> lockStream() override {
+    std::unique_lock<std::recursive_mutex> lockStream() override {
         lockStreamCalled = true;
         return AUBCommandStreamReceiver::AubFileStream::lockStream();
     }