1 files changed, 112 insertions, 0 deletions
diff --git a/fcp/client/opstats/opstats_logger_impl.h b/fcp/client/opstats/opstats_logger_impl.h
new file mode 100644
index 0000000..581406f
--- /dev/null
+++ b/fcp/client/opstats/opstats_logger_impl.h
@@ -0,0 +1,112 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef FCP_CLIENT_OPSTATS_OPSTATS_LOGGER_IMPL_H_
+#define FCP_CLIENT_OPSTATS_OPSTATS_LOGGER_IMPL_H_
+
+#include <string>
+
+#include "absl/synchronization/mutex.h"
+#include "fcp/client/flags.h"
+#include "fcp/client/log_manager.h"
+#include "fcp/client/opstats/opstats_db.h"
+#include "fcp/client/opstats/opstats_logger.h"
+#include "fcp/client/stats.h"
+#include "fcp/protos/opstats.pb.h"
+
+namespace fcp {
+namespace client {
+namespace opstats {
+
+// An implementation of OpStatsLogger backed by a database.
+class OpStatsLoggerImpl : public OpStatsLogger {
+ public:
+  // Creates a logger backed by an actual database. Populates the internal
+  // message with the provided session and population names.
+  OpStatsLoggerImpl(std::unique_ptr<OpStatsDb> db, LogManager* log_manager,
+                    const Flags* flags, const std::string& session_name,
+                    const std::string& population_name);
+
+  // Commits the cumulative message to the db.
+  ~OpStatsLoggerImpl() override;
+
+  // Adds an event and the given task name to the cumulative internal message,
+  // in a single transaction.
+  void AddEventAndSetTaskName(const std::string& task_name,
+                              OperationalStats::Event::EventKind event)
+      ABSL_LOCKS_EXCLUDED(mutex_) override;
+
+  // Adds an event to the cumulative internal message.
+  void AddEvent(OperationalStats::Event::EventKind event)
+      ABSL_LOCKS_EXCLUDED(mutex_) override;
+
+  // Adds an event and corresponding error message to the cumulative internal
+  // message.
+  void AddEventWithErrorMessage(OperationalStats::Event::EventKind event,
+                                const std::string& error_message) override;
+
+  // Updates info associated with a dataset created for a given collection in
+  // the cumulative internal message. If this is called multiple times for the
+  // same collection, the example counts and sizes will be aggregated in the
+  // underlying submessage.
+  void UpdateDatasetStats(const std::string& collection_uri,
+                          int additional_example_count,
+                          int64_t additional_example_size_bytes)
+      ABSL_LOCKS_EXCLUDED(mutex_) override;
+
+  // Adds network stats, replacing any old stats for the run, to the cumulative
+  // internal message.
+  void SetNetworkStats(const NetworkStats& network_stats)
+      ABSL_LOCKS_EXCLUDED(mutex_) override;
+
+  // Sets the retry window, replacing any old retry window, in the cumulative
+  // internal message. Any retry token in the retry window message is dropped.
+  void SetRetryWindow(
+      google::internal::federatedml::v2::RetryWindow retry_window)
+      ABSL_LOCKS_EXCLUDED(mutex_) override;
+
+  // Get the underlying opstats database.
+  OpStatsDb* GetOpStatsDb() override { return db_.get(); }
+
+  // Whether opstats is enabled. An instance of this class should only ever be
+  // created when opstats is enabled.
+  bool IsOpStatsEnabled() const override { return true; }
+
+  // Syncs all logged events to storage.
+  absl::Status CommitToStorage() override;
+
+  // Returns the task name of the currently executing task. Only returns a valid
+  // task name if called after `AddEventAndSetTaskName` is called.
+  std::string GetCurrentTaskName() override;
+
+ private:
+  // Helper for adding a new event of the specified kind to the cumulative
+  // message being stored in this class.
+  void AddNewEventToStats(OperationalStats::Event::EventKind kind)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Cumulative message storing information about this run.
+  OperationalStats stats_ ABSL_GUARDED_BY(mutex_);
+  bool already_committed_ ABSL_GUARDED_BY(mutex_) = false;
+  std::unique_ptr<OpStatsDb> db_;
+  LogManager* log_manager_;
+  absl::Mutex mutex_;
+};
+
+}  // namespace opstats
+}  // namespace client
+}  // namespace fcp
+
+#endif  // FCP_CLIENT_OPSTATS_OPSTATS_LOGGER_IMPL_H_