1 /* 2 * Copyright 2021 Google LLC 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef FCP_CLIENT_OPSTATS_OPSTATS_LOGGER_IMPL_H_ 17 #define FCP_CLIENT_OPSTATS_OPSTATS_LOGGER_IMPL_H_ 18 19 #include <string> 20 21 #include "absl/synchronization/mutex.h" 22 #include "fcp/client/flags.h" 23 #include "fcp/client/log_manager.h" 24 #include "fcp/client/opstats/opstats_db.h" 25 #include "fcp/client/opstats/opstats_logger.h" 26 #include "fcp/client/stats.h" 27 #include "fcp/protos/opstats.pb.h" 28 29 namespace fcp { 30 namespace client { 31 namespace opstats { 32 33 // An implementation of OpStatsLogger backed by a database. 34 class OpStatsLoggerImpl : public OpStatsLogger { 35 public: 36 // Creates a logger backed by an actual database. Populates the internal 37 // message with the provided session and population names. 38 OpStatsLoggerImpl(std::unique_ptr<OpStatsDb> db, LogManager* log_manager, 39 const Flags* flags, const std::string& session_name, 40 const std::string& population_name); 41 42 // Commits the cumulative message to the db. 43 ~OpStatsLoggerImpl() override; 44 45 // Adds an event and the given task name to the cumulative internal message, 46 // in a single transaction. 47 void AddEventAndSetTaskName(const std::string& task_name, 48 OperationalStats::Event::EventKind event) 49 ABSL_LOCKS_EXCLUDED(mutex_) override; 50 51 // Adds an event to the cumulative internal message. 52 void AddEvent(OperationalStats::Event::EventKind event) 53 ABSL_LOCKS_EXCLUDED(mutex_) override; 54 55 // Adds an event and corresponding error message to the cumulative internal 56 // message. 57 void AddEventWithErrorMessage(OperationalStats::Event::EventKind event, 58 const std::string& error_message) override; 59 60 // Updates info associated with a dataset created for a given collection in 61 // the cumulative internal message. If this is called multiple times for the 62 // same collection, the example counts and sizes will be aggregated in the 63 // underlying submessage. 64 void UpdateDatasetStats(const std::string& collection_uri, 65 int additional_example_count, 66 int64_t additional_example_size_bytes) 67 ABSL_LOCKS_EXCLUDED(mutex_) override; 68 69 // Adds network stats, replacing any old stats for the run, to the cumulative 70 // internal message. 71 void SetNetworkStats(const NetworkStats& network_stats) 72 ABSL_LOCKS_EXCLUDED(mutex_) override; 73 74 // Sets the retry window, replacing any old retry window, in the cumulative 75 // internal message. Any retry token in the retry window message is dropped. 76 void SetRetryWindow( 77 google::internal::federatedml::v2::RetryWindow retry_window) 78 ABSL_LOCKS_EXCLUDED(mutex_) override; 79 80 // Get the underlying opstats database. GetOpStatsDb()81 OpStatsDb* GetOpStatsDb() override { return db_.get(); } 82 83 // Whether opstats is enabled. An instance of this class should only ever be 84 // created when opstats is enabled. IsOpStatsEnabled()85 bool IsOpStatsEnabled() const override { return true; } 86 87 // Syncs all logged events to storage. 88 absl::Status CommitToStorage() override; 89 90 // Returns the task name of the currently executing task. Only returns a valid 91 // task name if called after `AddEventAndSetTaskName` is called. 92 std::string GetCurrentTaskName() override; 93 94 private: 95 // Helper for adding a new event of the specified kind to the cumulative 96 // message being stored in this class. 97 void AddNewEventToStats(OperationalStats::Event::EventKind kind) 98 ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_); 99 100 // Cumulative message storing information about this run. 101 OperationalStats stats_ ABSL_GUARDED_BY(mutex_); 102 bool already_committed_ ABSL_GUARDED_BY(mutex_) = false; 103 std::unique_ptr<OpStatsDb> db_; 104 LogManager* log_manager_; 105 absl::Mutex mutex_; 106 }; 107 108 } // namespace opstats 109 } // namespace client 110 } // namespace fcp 111 112 #endif // FCP_CLIENT_OPSTATS_OPSTATS_LOGGER_IMPL_H_ 113