From 21e3002d5a6e52c61286cf1ee67a477df75f82e5 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 2 Jun 2021 11:51:15 +1000
Subject: [PATCH] Support multiple output sinks in logging_resource_adaptor
 (#791)

Adds a constructor to `rmm::mr::logging_resource_adaptor` that takes an init list of spdlog sinks.  This enables creation of multiple sinks so that a `logging_resource_adaptor` may output to multiple places (For example to a file and an output stream). Also adds a test.

Closes #679

Authors:
  - Mark Harris (https://github.com/harrism)

Approvers:
  - Conor Hoekstra (https://github.com/codereport)
  - Rong Ou (https://github.com/rongou)

URL: https://github.com/rapidsai/rmm/pull/791
---
 .../mr/device/logging_resource_adaptor.hpp    | 11 +++++++
 tests/logger_tests.cpp                        | 30 +++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/include/rmm/mr/device/logging_resource_adaptor.hpp b/include/rmm/mr/device/logging_resource_adaptor.hpp
index 757d5c762..bd5c38685 100644
--- a/include/rmm/mr/device/logging_resource_adaptor.hpp
+++ b/include/rmm/mr/device/logging_resource_adaptor.hpp
@@ -15,6 +15,7 @@
  */
 #pragma once
 
+#include <spdlog/common.h>
 #include <rmm/mr/device/device_memory_resource.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
@@ -107,6 +108,16 @@ class logging_resource_adaptor final : public device_memory_resource {
     init_logger(auto_flush);
   }
 
+  logging_resource_adaptor(Upstream* upstream,
+                           spdlog::sinks_init_list sinks,
+                           bool auto_flush = false)
+    : logger_{std::make_shared<spdlog::logger>("RMM", sinks)}, upstream_{upstream}
+  {
+    RMM_EXPECTS(nullptr != upstream, "Unexpected null upstream resource pointer.");
+
+    init_logger(auto_flush);
+  }
+
   logging_resource_adaptor()                                = delete;
   ~logging_resource_adaptor()                               = default;
   logging_resource_adaptor(logging_resource_adaptor const&) = delete;
diff --git a/tests/logger_tests.cpp b/tests/logger_tests.cpp
index dfa461850..b343c7e35 100644
--- a/tests/logger_tests.cpp
+++ b/tests/logger_tests.cpp
@@ -105,6 +105,36 @@ TEST(Adaptor, FilenameConstructor)
   expect_log_events(filename, expected_events);
 }
 
+TEST(Adaptor, MultiSinkConstructor)
+{
+  std::string filename1{"logs/test_multi_1.txt"};
+  std::string filename2{"logs/test_multi_2.txt"};
+  rmm::mr::cuda_memory_resource upstream;
+
+  auto file_sink1 = std::make_shared<spdlog::sinks::basic_file_sink_mt>(filename1, true);
+  auto file_sink2 = std::make_shared<spdlog::sinks::basic_file_sink_mt>(filename2, true);
+
+  rmm::mr::logging_resource_adaptor<rmm::mr::cuda_memory_resource> log_mr{&upstream,
+                                                                          {file_sink1, file_sink2}};
+
+  auto p0 = log_mr.allocate(100);
+  auto p1 = log_mr.allocate(42);
+  log_mr.deallocate(p0, 100);
+  log_mr.deallocate(p1, 42);
+  log_mr.flush();
+
+  using rmm::detail::action;
+  using rmm::detail::event;
+
+  std::vector<event> expected_events{{action::ALLOCATE, 100, p0},
+                                     {action::ALLOCATE, 42, p1},
+                                     {action::FREE, 100, p0},
+                                     {action::FREE, 42, p1}};
+
+  expect_log_events(filename1, expected_events);
+  expect_log_events(filename2, expected_events);
+}
+
 TEST(Adaptor, Factory)
 {
   std::string filename{"logs/test2.txt"};