1 files changed, 71 insertions, 0 deletions
diff --git a/test/util/save_util.cc b/test/util/save_util.cc
new file mode 100644
index 000000000..384d626f0
--- /dev/null
+++ b/test/util/save_util.cc
@@ -0,0 +1,71 @@
+// Copyright 2018 The gVisor Authors.
+//
+// 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.
+
+#include "test/util/save_util.h"
+
+#include <stddef.h>
+#include <stdlib.h>
+#include <unistd.h>
+
+#include <atomic>
+#include <cerrno>
+
+#define GVISOR_COOPERATIVE_SAVE_TEST "GVISOR_COOPERATIVE_SAVE_TEST"
+
+namespace gvisor {
+namespace testing {
+namespace {
+
+enum class CooperativeSaveMode {
+  kUnknown = 0,  // cooperative_save_mode is statically-initialized to 0
+  kAvailable,
+  kNotAvailable,
+};
+
+std::atomic<CooperativeSaveMode> cooperative_save_mode;
+
+bool CooperativeSaveEnabled() {
+  auto mode = cooperative_save_mode.load();
+  if (mode == CooperativeSaveMode::kUnknown) {
+    mode = (getenv(GVISOR_COOPERATIVE_SAVE_TEST) != nullptr)
+               ? CooperativeSaveMode::kAvailable
+               : CooperativeSaveMode::kNotAvailable;
+    cooperative_save_mode.store(mode);
+  }
+  return mode == CooperativeSaveMode::kAvailable;
+}
+
+std::atomic<int> save_disable;
+
+}  // namespace
+
+DisableSave::DisableSave() { save_disable++; }
+
+DisableSave::~DisableSave() { reset(); }
+
+void DisableSave::reset() {
+  if (!reset_) {
+    reset_ = true;
+    save_disable--;
+  }
+}
+
+namespace internal {
+bool ShouldSave() {
+  return CooperativeSaveEnabled() && (save_disable.load() == 0);
+}
+}  // namespace internal
+
+}  // namespace testing
+}  // namespace gvisor