Refactor dry_run and walk_tree into RunMode.

This is shared across the garage tools. This enabled me to fix the remaining TODO in OSTreeObject::CurlDone(). Signed-off-by: Patrick Vacek <[email protected]>
advancedtelematic · Dec 3, 2018 · 43e91da · 43e91da
1 parent fc0f8ea
commit 43e91da
Show file tree

Hide file tree

Showing 14 changed files with 95 additions and 69 deletions.
diff --git a/src/sota_tools/CMakeLists.txt b/src/sota_tools/CMakeLists.txt
@@ -93,6 +93,7 @@ set(ALL_SOTA_TOOLS_HEADERS
     accumulator.h
     authenticate.h
     deploy.h
+    garage_common.h
     oauth2.h
     ostree_dir_repo.h
     ostree_hash.h

diff --git a/src/sota_tools/deploy.cc b/src/sota_tools/deploy.cc
@@ -12,7 +12,7 @@
 #include "utilities/utils.h"
 
 bool UploadToTreehub(const OSTreeRepo::ptr &src_repo, const ServerCredentials &push_credentials,
-                     const OSTreeHash &ostree_commit, const std::string &cacerts, const bool dryrun,
+                     const OSTreeHash &ostree_commit, const std::string &cacerts, const RunMode mode,
                      const int max_curl_requests) {
   TreehubServer push_server;
   assert(max_curl_requests > 0);
@@ -29,7 +29,7 @@ bool UploadToTreehub(const OSTreeRepo::ptr &src_repo, const ServerCredentials &p
     return false;
   }
 
-  RequestPool request_pool(push_server, max_curl_requests);
+  RequestPool request_pool(push_server, max_curl_requests, mode);
 
   // Add commit object to the queue.
   request_pool.AddQuery(root_object);
@@ -39,11 +39,11 @@ bool UploadToTreehub(const OSTreeRepo::ptr &src_repo, const ServerCredentials &p
   // OSTreeObject::CurlDone() adds new requests to the pool and stops the pool
   // on error.
   do {
-    request_pool.Loop(dryrun);
+    request_pool.Loop();
   } while (root_object->is_on_server() != PresenceOnServer::kObjectPresent && !request_pool.is_stopped());
 
   if (root_object->is_on_server() == PresenceOnServer::kObjectPresent) {
-    if (!dryrun) {
+    if (mode == RunMode::kDefault) {
       LOG_INFO << "Upload to Treehub complete after " << request_pool.total_requests_made() << " requests";
     } else {
       LOG_INFO << "Dry run. No objects uploaded.";
@@ -100,7 +100,7 @@ bool OfflineSignRepo(const ServerCredentials &push_credentials, const std::strin
 }
 
 bool PushRootRef(const ServerCredentials &push_credentials, const OSTreeRef &ref, const std::string &cacerts,
-                 const bool dry_run) {
+                 const RunMode mode) {
   if (push_credentials.CanSignOffline()) {
     // In general, this is the wrong thing.  We should be using offline signing
     // if private key material is present in credentials.zip
@@ -114,7 +114,7 @@ bool PushRootRef(const ServerCredentials &push_credentials, const OSTreeRef &ref
     return false;
   }
 
-  if (!dry_run) {
+  if (mode == RunMode::kDefault) {
     CurlEasyWrapper easy_handle;
     curlEasySetoptWrapper(easy_handle.get(), CURLOPT_VERBOSE, get_curlopt_verbose());
     ref.PushRef(push_server, easy_handle.get());

diff --git a/src/sota_tools/deploy.h b/src/sota_tools/deploy.h
@@ -3,6 +3,7 @@
 
 #include <string>
 
+#include "garage_common.h"
 #include "ostree_ref.h"
 #include "ostree_repo.h"
 #include "server_credentials.h"
@@ -17,11 +18,11 @@
  * \param push_credentials
  * \param ostree_commit
  * \param cacerts
- * \param dryrun
+ * \param mode
  * \param max_curl_requests
  */
 bool UploadToTreehub(const OSTreeRepo::ptr& src_repo, const ServerCredentials& push_credentials,
-                     const OSTreeHash& ostree_commit, const std::string& cacerts, bool dryrun, int max_curl_requests);
+                     const OSTreeHash& ostree_commit, const std::string& cacerts, RunMode mode, int max_curl_requests);
 
 /**
  * Use the garage-sign tool and the images targets.json keys in credentials.zip
@@ -34,6 +35,6 @@ bool OfflineSignRepo(const ServerCredentials& push_credentials, const std::strin
  * Update images/targets.json by pushing the OSTree commit hash to /refs/heads/qemux86-64
  */
 bool PushRootRef(const ServerCredentials& push_credentials, const OSTreeRef& ref, const std::string& cacerts,
-                 bool dry_run);
+                 RunMode mode);
 
 #endif
diff --git a/src/sota_tools/deploy_test.cc b/src/sota_tools/deploy_test.cc
@@ -2,6 +2,7 @@
 
 #include "crypto/crypto.h"
 #include "deploy.h"
+#include "garage_common.h"
 #include "ostree_dir_repo.h"
 #include "ostree_http_repo.h"
 #include "ostree_ref.h"
@@ -20,7 +21,7 @@ TEST(deploy, UploadToTreehub) {
   const uint8_t hash[32] = {0x16, 0xef, 0x2f, 0x26, 0x29, 0xdc, 0x92, 0x63, 0xfd, 0xf3, 0xc0,
                             0xf0, 0x32, 0x56, 0x3a, 0x2d, 0x75, 0x76, 0x23, 0xbb, 0xc1, 0x1c,
                             0xf9, 0x9d, 0xf2, 0x5c, 0x3c, 0x3f, 0x25, 0x8d, 0xcc, 0xbe};
-  UploadToTreehub(src_repo, ServerCredentials(filepath), OSTreeHash(hash), cert_path.string(), false, 2);
+  UploadToTreehub(src_repo, ServerCredentials(filepath), OSTreeHash(hash), cert_path.string(), RunMode::kDefault, 2);
 
   int result = system(
       (std::string("diff -r ") + (temp_dir.Path() / "objects/").string() + " tests/sota_tools/repo/objects/").c_str());

diff --git a/src/sota_tools/garage_check.cc b/src/sota_tools/garage_check.cc
@@ -9,6 +9,7 @@
 
 #include "accumulator.h"
 #include "authenticate.h"
+#include "garage_common.h"
 #include "logging/logging.h"
 #include "ostree_http_repo.h"
 #include "ostree_object.h"
@@ -38,7 +39,7 @@ int main(int argc, char **argv) {
   string cacerts;
 
   int verbosity;
-  bool walk_tree = false;
+  RunMode mode = RunMode::kDefault;
   po::options_description desc("garage-check command line options");
   // clang-format off
   desc.add_options()
@@ -87,7 +88,7 @@ int main(int argc, char **argv) {
   }
 
   if (vm.count("walk-tree") != 0u) {
-    walk_tree = true;
+    mode = RunMode::kWalkTree;
   }
 
   TreehubServer treehub;
@@ -128,7 +129,7 @@ int main(int argc, char **argv) {
   bool is_commit = true;
   curl_easy_getinfo(curl.get(), CURLINFO_RESPONSE_CODE, &http_code);
   if (http_code == 404) {
-    if (!walk_tree) {
+    if (mode != RunMode::kWalkTree) {
       LOG_FATAL << "OSTree commit " << ref << " is missing in treehub";
       return EXIT_FAILURE;
     } else {
@@ -138,17 +139,17 @@ int main(int argc, char **argv) {
     LOG_FATAL << "Error " << http_code << " getting OSTree ref " << ref << " from treehub";
     return EXIT_FAILURE;
   }
-  if (!walk_tree) {
+  if (mode != RunMode::kWalkTree) {
     LOG_INFO << "OSTree commit " << ref << " is found on treehub";
   }
 
-  if (walk_tree) {
+  if (mode == RunMode::kWalkTree) {
     // Walk the entire tree and check for all objects.
     OSTreeHttpRepo dest_repo(&treehub);
     OSTreeHash hash = OSTreeHash::Parse(ref);
     OSTreeObject::ptr input_object = dest_repo.GetObject(hash);
 
-    RequestPool request_pool(treehub, 1);
+    RequestPool request_pool(treehub, 1, mode);
 
     // Add input object to the queue.
     request_pool.AddQuery(input_object);
@@ -157,10 +158,9 @@ int main(int argc, char **argv) {
     // request_pool takes care of holding number of outstanding requests below.
     // OSTreeObject::CurlDone() adds new requests to the pool and stops the pool
     // on error.
-    const bool dryrun = true;
     do {
-      request_pool.Loop(dryrun);
-    } while (!request_pool.is_stopped()); // TODO: we probably need a better stop condition!
+      request_pool.Loop();
+    } while (!request_pool.is_stopped());  // TODO: this will probably run forever!
 
     if (input_object->is_on_server() == PresenceOnServer::kObjectPresent) {
       LOG_INFO << "Dry run. No objects uploaded.";

diff --git a/src/sota_tools/garage_common.h b/src/sota_tools/garage_common.h
@@ -0,0 +1,15 @@
+#ifndef GARAGE_COMMON_H_
+#define GARAGE_COMMON_H_
+
+/** Execution mode to run garage tools in. */
+enum class RunMode {
+  /** Default operation. Upload objects to server if necessary and relevant. */
+  kDefault = 0,
+  /** Dry run. Do not upload any objects. */
+  kDryRun,
+  /** Walk the entire tree (without uploading). Do not assume that if an object
+   * exists, its parents must also exist. */
+  kWalkTree,
+};
+
+#endif  // GARAGE_COMMON_H_
diff --git a/src/sota_tools/garage_deploy.cc b/src/sota_tools/garage_deploy.cc
@@ -6,6 +6,7 @@
 #include "accumulator.h"
 #include "authenticate.h"
 #include "deploy.h"
+#include "garage_common.h"
 #include "logging/logging.h"
 #include "ostree_http_repo.h"
 
@@ -20,7 +21,7 @@ int main(int argc, char **argv) {
   boost::filesystem::path push_cred;
   std::string hardwareids;
   std::string cacerts;
-  bool dry_run = false;
+  RunMode mode = RunMode::kDefault;
   po::options_description desc("garage-deploy command line options");
   // clang-format off
   desc.add_options()
@@ -76,7 +77,7 @@ int main(int argc, char **argv) {
   }
 
   if (vm.count("dry-run") != 0u) {
-    dry_run = true;
+    mode = RunMode::kDryRun;
   }
 
   ServerCredentials push_credentials(push_cred);
@@ -93,12 +94,12 @@ int main(int argc, char **argv) {
     OSTreeHash commit(OSTreeHash::Parse(ostree_commit));
     // Since the fetches happen on a single thread in OSTreeHttpRepo, there
     // isn't really any reason to upload in parallel
-    if (!UploadToTreehub(src_repo, push_credentials, commit, cacerts, dry_run, 1)) {
+    if (!UploadToTreehub(src_repo, push_credentials, commit, cacerts, mode, 1)) {
       LOG_FATAL << "Upload to treehub failed";
       return EXIT_FAILURE;
     }
 
-    if (!dry_run) {
+    if (mode == RunMode::kDefault) {
       if (push_credentials.CanSignOffline()) {
         bool ok = OfflineSignRepo(ServerCredentials(push_credentials.GetPathOnDisk()), name, commit, hardwareids);
         return static_cast<int>(!ok);

diff --git a/src/sota_tools/garage_push.cc b/src/sota_tools/garage_push.cc
@@ -6,6 +6,7 @@
 
 #include "accumulator.h"
 #include "deploy.h"
+#include "garage_common.h"
 #include "logging/logging.h"
 #include "ostree_dir_repo.h"
 #include "ostree_repo.h"
@@ -23,7 +24,7 @@ int main(int argc, char **argv) {
   boost::filesystem::path credentials_path;
 
   int verbosity;
-  bool dry_run = false;
+  RunMode mode = RunMode::kDefault;
   po::options_description desc("garage-push command line options");
   // clang-format off
   desc.add_options()
@@ -81,7 +82,7 @@ int main(int argc, char **argv) {
   }
 
   if (vm.count("dry-run") != 0u) {
-    dry_run = true;
+    mode = RunMode::kDryRun;
   }
 
   if (max_curl_requests < 1) {
@@ -105,7 +106,7 @@ int main(int argc, char **argv) {
     }
 
     OSTreeHash commit(ostree_ref.GetHash());
-    bool ok = UploadToTreehub(src_repo, push_credentials, commit, cacerts, dry_run, max_curl_requests);
+    bool ok = UploadToTreehub(src_repo, push_credentials, commit, cacerts, mode, max_curl_requests);
 
     if (!ok) {
       LOG_FATAL << "Upload to treehub failed";
@@ -115,7 +116,7 @@ int main(int argc, char **argv) {
     if (push_credentials.CanSignOffline()) {
       LOG_INFO << "Credentials contain offline signing keys, not pushing root ref";
     } else {
-      ok = PushRootRef(push_credentials, ostree_ref, cacerts, dry_run);
+      ok = PushRootRef(push_credentials, ostree_ref, cacerts, mode);
       if (!ok) {
         LOG_FATAL << "Could not push root reference to treehub";
         return EXIT_FAILURE;

diff --git a/src/sota_tools/ostree_http_repo_test.cc b/src/sota_tools/ostree_http_repo_test.cc
@@ -1,6 +1,7 @@
 #include <gtest/gtest.h>
 
 #include "deploy.h"
+#include "garage_common.h"
 #include "ostree_http_repo.h"
 #include "ostree_ref.h"
 #include "server_credentials.h"
@@ -89,7 +90,7 @@ TEST(http_repo, bad_connection) {
   const uint8_t hash[32] = {0x16, 0xef, 0x2f, 0x26, 0x29, 0xdc, 0x92, 0x63, 0xfd, 0xf3, 0xc0,
                             0xf0, 0x32, 0x56, 0x3a, 0x2d, 0x75, 0x76, 0x23, 0xbb, 0xc1, 0x1c,
                             0xf9, 0x9d, 0xf2, 0x5c, 0x3c, 0x3f, 0x25, 0x8d, 0xcc, 0xbe};
-  UploadToTreehub(src_repo, ServerCredentials(filepath), OSTreeHash(hash), cert_path.string(), false, 1);
+  UploadToTreehub(src_repo, ServerCredentials(filepath), OSTreeHash(hash), cert_path.string(), RunMode::kDefault, 1);
 
   int result = system(
       (std::string("diff -r ") + (src_repo->root() / "objects/").string() + " tests/sota_tools/repo/objects/").c_str());

diff --git a/src/sota_tools/ostree_object.cc b/src/sota_tools/ostree_object.cc
@@ -198,8 +198,8 @@ void OSTreeObject::MakeTestRequest(const TreehubServer &push_target, CURLM *curl
   request_start_time_ = std::chrono::steady_clock::now();
 }
 
-void OSTreeObject::Upload(const TreehubServer &push_target, CURLM *curl_multi_handle, const bool dryrun) {
-  if (!dryrun) {
+void OSTreeObject::Upload(const TreehubServer &push_target, CURLM *curl_multi_handle, const RunMode mode) {
+  if (mode == RunMode::kDefault) {
     LOG_INFO << "Uploading " << object_name_;
   } else {
     LOG_INFO << "Would upload " << object_name_;
@@ -241,6 +241,21 @@ void OSTreeObject::Upload(const TreehubServer &push_target, CURLM *curl_multi_ha
   request_start_time_ = std::chrono::steady_clock::now();
 }
 
+void OSTreeObject::CheckChildren(RequestPool &pool) {
+  try {
+    PopulateChildren();
+    LOG_TRACE << "Children of " << object_name_ << ": " << children_.size();
+    if (children_ready()) {
+      pool.AddUpload(this);
+    } else {
+      QueryChildren(pool);
+    }
+  } catch (const OSTreeObjectMissing &error) {
+    LOG_ERROR << "Source OSTree repo does not contain object " << error.missing_object();
+    pool.Abort();
+  }
+}
+
 void OSTreeObject::PresenceError(RequestPool &pool, const int64_t rescode) {
   is_on_server_ = PresenceOnServer::kObjectStateUnknown;
   LOG_WARNING << "OSTree query reported an error code: " << rescode << " retrying...";
@@ -276,34 +291,15 @@ void OSTreeObject::CurlDone(CURLM *curl_multi_handle, RequestPool &pool) {
       LOG_INFO << "Already present: " << object_name_;
       is_on_server_ = PresenceOnServer::kObjectPresent;
       last_operation_result_ = ServerResponse::kOk;
-      // TODO: Fix
-      // NotifyParents(pool);
-      try {
-        PopulateChildren();
-        LOG_INFO << "Children: " << children_.size();
-        if (children_ready()) {
-          pool.AddUpload(this);
-        } else {
-          QueryChildren(pool);
-        }
-      } catch (const OSTreeObjectMissing &error) {
-        LOG_ERROR << "Source OSTree repo does not contain object " << error.missing_object();
-        pool.Abort();
+      if (pool.run_mode() != RunMode::kWalkTree) {
+        NotifyParents(pool);
+      } else {
+        CheckChildren(pool);
       }
     } else if (rescode == 404) {
       is_on_server_ = PresenceOnServer::kObjectMissing;
       last_operation_result_ = ServerResponse::kOk;
-      try {
-        PopulateChildren();
-        if (children_ready()) {
-          pool.AddUpload(this);
-        } else {
-          QueryChildren(pool);
-        }
-      } catch (const OSTreeObjectMissing &error) {
-        LOG_ERROR << "Source OSTree repo does not contain object " << error.missing_object();
-        pool.Abort();
-      }
+      CheckChildren(pool);
     } else {
       PresenceError(pool, rescode);
     }

diff --git a/src/sota_tools/ostree_object.h b/src/sota_tools/ostree_object.h
@@ -10,6 +10,7 @@
 #include <boost/filesystem.hpp>
 #include <boost/intrusive_ptr.hpp>
 
+#include "garage_common.h"
 #include "treehub_server.h"
 
 class OSTreeRepo;
@@ -44,7 +45,7 @@ class OSTreeObject {
   void MakeTestRequest(const TreehubServer& push_target, CURLM* curl_multi_handle);
 
   /* Upload this object to the destination server. */
-  void Upload(const TreehubServer& push_target, CURLM* curl_multi_handle, bool dryrun);
+  void Upload(const TreehubServer& push_target, CURLM* curl_multi_handle, RunMode mode);
 
   /* Process a completed curl transaction (presence check or upload). */
   void CurlDone(CURLM* curl_multi_handle, RequestPool& pool);
@@ -79,6 +80,10 @@ class OSTreeObject {
 
   std::string Url() const;
 
+  /* Check for children. If they are all present, upload this object. Otherwise,
+   * query each of the children. */
+  void CheckChildren(RequestPool& pool);
+
   /* Handle an error from a presence check. */
   void PresenceError(RequestPool& pool, int64_t rescode);