4 files changed, 118 insertions, 103 deletions

diff --git a/absl/flags/internal/commandlineflag.cc b/absl/flags/internal/commandlineflag.cc
index 53e2b84ecc13..99f73611d8f1 100644
--- a/absl/flags/internal/commandlineflag.cc
+++ b/absl/flags/internal/commandlineflag.cc
@@ -149,6 +149,12 @@ std::string CommandLineFlag::CurrentValue() const {
   return Unparse(marshalling_op_, cur_);
 }
 
+int64_t CommandLineFlag::MutationCounter() const {
+  absl::MutexLock l(InitFlagIfNecessary());
+
+  return counter_;
+}
+
 // Attempts to parse supplied `value` string using parsing routine in the `flag`
 // argument. If parsing is successful, it will try to validate that the parsed
 // value is valid for the specified 'flag'. Finally this function stores the
diff --git a/absl/flags/internal/commandlineflag.h b/absl/flags/internal/commandlineflag.h
index 284286b6be18..528d3106bfa8 100644
--- a/absl/flags/internal/commandlineflag.h
+++ b/absl/flags/internal/commandlineflag.h
@@ -17,6 +17,7 @@
 #define ABSL_FLAGS_INTERNAL_COMMANDLINEFLAG_H_
 
 #include <atomic>
+#include <memory>
 
 #include "absl/base/macros.h"
 #include "absl/flags/marshalling.h"
@@ -186,6 +187,16 @@ class HelpText {
   const char* help_message_;
 };
 
+// Handle to FlagState objects. Specific flag state objects will restore state
+// of a flag produced this flag state from method CommandLineFlag::SaveState().
+class FlagStateInterface {
+ public:
+  virtual ~FlagStateInterface() {}
+
+  // Restores the flag originated this object to the saved state.
+  virtual void Restore() const = 0;
+};
+
 // Holds all information for a flag.
 class CommandLineFlag {
  public:
@@ -239,7 +250,6 @@ class CommandLineFlag {
   virtual void StoreAtomic() {}
 
   // Interfaces to operate on validators.
-  virtual bool HasValidatorFn() const { return false; }
   virtual bool InvokeValidator(const void* /*value*/) const { return true; }
   // Invoke the flag validators for old flags.
   // TODO(rogeeff): implement proper validators for Abseil Flags
@@ -264,6 +274,10 @@ class CommandLineFlag {
     return res;
   }
 
+  // Interface to save flag to some persistent state. Returns current flag state
+  // or nullptr if flag does not support saving and restoring a state.
+  virtual std::unique_ptr<FlagStateInterface> SaveState() = 0;
+
   // Interfaces to overate on callbacks.
   virtual void InvokeCallback() {}
 
@@ -285,6 +299,9 @@ class CommandLineFlag {
  protected:
   ~CommandLineFlag() = default;
 
+  // Thread safe access to mutation counter.
+  int64_t MutationCounter() const;
+
   const char* const name_;
   const HelpText help_;
   const char* const filename_;
@@ -323,11 +340,6 @@ class CommandLineFlag {
   friend bool TryParseLocked(CommandLineFlag* flag, void* dst,
                              absl::string_view value, std::string* err);
   friend absl::Mutex* InitFlag(CommandLineFlag* flag);
-
-  // This is a short term, until we completely rework persistent state
-  // storage API.
-  virtual void* GetValidator() const { return nullptr; }
-  virtual bool SetValidator(void*) { return false; }
 };
 
 // Update any copy of the flag value that is stored in an atomic word.
diff --git a/absl/flags/internal/flag.h b/absl/flags/internal/flag.h
index 16330380f548..2b21c4407363 100644
--- a/absl/flags/internal/flag.h
+++ b/absl/flags/internal/flag.h
@@ -20,12 +20,44 @@
 
 #include "absl/flags/internal/commandlineflag.h"
 #include "absl/flags/internal/registry.h"
+#include "absl/memory/memory.h"
+#include "absl/strings/str_cat.h"
 
 namespace absl {
 namespace flags_internal {
 
 constexpr int64_t AtomicInit() { return 0xababababababababll; }
 
+template <typename T>
+class Flag;
+
+template <typename T>
+class FlagState : public flags_internal::FlagStateInterface {
+ public:
+  FlagState(Flag<T>* flag, T&& cur, bool modified, bool on_command_line,
+            int64_t counter)
+      : flag_(flag),
+        cur_value_(std::move(cur)),
+        modified_(modified),
+        on_command_line_(on_command_line),
+        counter_(counter) {}
+
+  ~FlagState() override = default;
+
+ private:
+  friend class Flag<T>;
+
+  // Restores the flag to the saved state.
+  void Restore() const override;
+
+  // Flag and saved flag data.
+  Flag<T>* flag_;
+  T cur_value_;
+  bool modified_;
+  bool on_command_line_;
+  int64_t counter_;
+};
+
 // Signature for the mutation callback used by watched Flags
 // The callback is noexcept.
 // TODO(rogeeff): add noexcept after C++17 support is added.
@@ -98,15 +130,12 @@ class Flag final : public flags_internal::CommandLineFlag {
 
     InvokeCallback();
   }
-  void InvokeCallback() override
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(locks_->primary_mu) {
-    flags_internal::InvokeCallback(&locks_->primary_mu, &locks_->callback_mu,
-                                   callback_);
-  }
 
  private:
+  friend class FlagState<T>;
+
   void Destroy() const override {
-    // Values are heap allocated Abseil Flags.
+    // Values are heap allocated for Abseil Flags.
     if (cur_) Delete(op_, cur_);
     if (def_) Delete(op_, def_);
 
@@ -121,6 +150,41 @@ class Flag final : public flags_internal::CommandLineFlag {
     }
   }
 
+  // Interfaces to save and restore flags to/from persistent state.
+  // Returns current flag state or nullptr if flag does not support
+  // saving and restoring a state.
+  std::unique_ptr<flags_internal::FlagStateInterface> SaveState() override {
+    T curr_value = Get();
+
+    absl::MutexLock l(InitFlagIfNecessary());
+
+    return absl::make_unique<flags_internal::FlagState<T>>(
+        this, std::move(curr_value), modified_, on_command_line_, counter_);
+  }
+
+  // Restores the flag state to the supplied state object. If there is
+  // nothing to restore returns false. Otherwise returns true.
+  bool RestoreState(const flags_internal::FlagState<T>& flag_state) {
+    if (MutationCounter() == flag_state.counter_) return false;
+
+    Set(flag_state.cur_value_);
+
+    // Race condition here? This should disappear once we move the rest of the
+    // flag's data into Flag's internals.
+
+    absl::MutexLock l(InitFlagIfNecessary());
+    modified_ = flag_state.modified_;
+    on_command_line_ = flag_state.on_command_line_;
+    return true;
+  }
+
+  // Interfaces to overate on callbacks.
+  void InvokeCallback() override
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(locks_->primary_mu) {
+    flags_internal::InvokeCallback(&locks_->primary_mu, &locks_->callback_mu,
+                                   callback_);
+  }
+
   // Flag's data
   // For some types, a copy of the current value is kept in an atomically
   // accessible field.
@@ -128,6 +192,15 @@ class Flag final : public flags_internal::CommandLineFlag {
   FlagCallback callback_;  // Mutation callback
 };
 
+template <typename T>
+inline void FlagState<T>::Restore() const {
+  if (flag_->RestoreState(*this)) {
+    ABSL_INTERNAL_LOG(INFO,
+                      absl::StrCat("Restore saved value of ", flag_->Name(),
+                                   " to: ", flag_->CurrentValue()));
+  }
+}
+
 // This class facilitates Flag object registration and tail expression-based
 // flag definition, for example:
 // ABSL_FLAG(int, foo, 42, "Foo help").OnUpdate(NotifyFooWatcher);
diff --git a/absl/flags/internal/registry.cc b/absl/flags/internal/registry.cc
index 4bea313bb1eb..6b2564d13a78 100644
--- a/absl/flags/internal/registry.cc
+++ b/absl/flags/internal/registry.cc
@@ -188,114 +188,34 @@ CommandLineFlag* FlagRegistry::FindRetiredFlagLocked(absl::string_view name) {
 
 class FlagSaverImpl {
  public:
-  // Constructs an empty FlagSaverImpl object.
-  FlagSaverImpl() {}
-  ~FlagSaverImpl() {
-    // reclaim memory from each of our CommandLineFlags
-    for (const SavedFlag& src : backup_registry_) {
-      Delete(src.op, src.current);
-      Delete(src.op, src.default_value);
-    }
-  }
+  FlagSaverImpl() = default;
+  FlagSaverImpl(const FlagSaverImpl&) = delete;
+  void operator=(const FlagSaverImpl&) = delete;
 
   // Saves the flag states from the flag registry into this object.
   // It's an error to call this more than once.
-  // Must be called when the registry mutex is not held.
   void SaveFromRegistry() {
     assert(backup_registry_.empty());  // call only once!
-    SavedFlag saved;
     flags_internal::ForEachFlag([&](flags_internal::CommandLineFlag* flag) {
-      if (flag->IsRetired()) return;
-
-      saved.name = flag->Name();
-      saved.op = flag->op_;
-      saved.marshalling_op = flag->marshalling_op_;
-      {
-        absl::MutexLock l(flag->InitFlagIfNecessary());
-        saved.validator = flag->GetValidator();
-        saved.modified = flag->modified_;
-        saved.on_command_line = flag->on_command_line_;
-        saved.current = Clone(saved.op, flag->cur_);
-        saved.default_value = Clone(saved.op, flag->def_);
-        saved.counter = flag->counter_;
+      if (auto flag_state = flag->SaveState()) {
+        backup_registry_.emplace_back(std::move(flag_state));
       }
-      backup_registry_.push_back(saved);
     });
   }
 
-  // Restores the saved flag states into the flag registry.  We
-  // assume no flags were added or deleted from the registry since
-  // the SaveFromRegistry; if they were, that's trouble!  Must be
-  // called when the registry mutex is not held.
+  // Restores the saved flag states into the flag registry.
   void RestoreToRegistry() {
-    FlagRegistry* const global_registry = FlagRegistry::GlobalRegistry();
-    FlagRegistryLock frl(global_registry);
-    for (const SavedFlag& src : backup_registry_) {
-      CommandLineFlag* flag = global_registry->FindFlagLocked(src.name);
-      // If null, flag got deleted from registry.
-      if (!flag) continue;
-
-      bool restored = false;
-      {
-        // This function encapsulate the lock.
-        flag->SetValidator(src.validator);
-
-        absl::MutexLock l(flag->InitFlagIfNecessary());
-        flag->modified_ = src.modified;
-        flag->on_command_line_ = src.on_command_line;
-        if (flag->counter_ != src.counter ||
-            ChangedDirectly(flag, src.default_value, flag->def_)) {
-          restored = true;
-          Copy(src.op, src.default_value, flag->def_);
-        }
-        if (flag->counter_ != src.counter ||
-            ChangedDirectly(flag, src.current, flag->cur_)) {
-          restored = true;
-          Copy(src.op, src.current, flag->cur_);
-          UpdateCopy(flag);
-          flag->InvokeCallback();
-        }
-      }
-
-      if (restored) {
-        flag->counter_++;
-
-        // Revalidate the flag because the validator might store state based
-        // on the flag's value, which just changed due to the restore.
-        // Failing validation is ignored because it's assumed that the flag
-        // was valid previously and there's little that can be done about it
-        // here, anyway.
-        flag->ValidateInputValue(flag->CurrentValue());
-
-        ABSL_INTERNAL_LOG(
-            INFO, absl::StrCat("Restore saved value of ", flag->Name(), ": ",
-                               Unparse(src.marshalling_op, src.current)));
-      }
+    for (const auto& flag_state : backup_registry_) {
+      flag_state->Restore();
     }
   }
 
  private:
-  struct SavedFlag {
-    absl::string_view name;
-    FlagOpFn op;
-    FlagMarshallingOpFn marshalling_op;
-    int64_t counter;
-    void* validator;
-    bool modified;
-    bool on_command_line;
-    const void* current;        // nullptr after restore
-    const void* default_value;  // nullptr after restore
-  };
-
-  std::vector<SavedFlag> backup_registry_;
-
-  FlagSaverImpl(const FlagSaverImpl&);  // no copying!
-  void operator=(const FlagSaverImpl&);
+  std::vector<std::unique_ptr<flags_internal::FlagStateInterface>>
+      backup_registry_;
 };
 
-FlagSaver::FlagSaver() : impl_(new FlagSaverImpl()) {
-  impl_->SaveFromRegistry();
-}
+FlagSaver::FlagSaver() : impl_(new FlagSaverImpl) { impl_->SaveFromRegistry(); }
 
 void FlagSaver::Ignore() {
   delete impl_;
@@ -376,6 +296,10 @@ class RetiredFlagObj final : public flags_internal::CommandLineFlag {
 
     delete this;
   }
+
+  std::unique_ptr<flags_internal::FlagStateInterface> SaveState() override {
+    return nullptr;
+  }
 };
 
 }  // namespace