1 files changed, 59 insertions, 4 deletions

diff --git a/src/main/java/com/googlesource/gerrit/plugins/replication/StartCommand.java b/src/main/java/com/googlesource/gerrit/plugins/replication/StartCommand.java
index 5bcb7cb..bcdc3d9 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/replication/StartCommand.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/replication/StartCommand.java
@@ -24,20 +24,30 @@ import com.google.inject.Inject;
 
 import org.kohsuke.args4j.Argument;
 import org.kohsuke.args4j.Option;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 
 @RequiresCapability(GlobalCapability.START_REPLICATION)
 @CommandMetaData(name="start", descr="Start replication for specific project or all projects")
 final class StartCommand extends SshCommand {
+  private static final Logger log = LoggerFactory.getLogger(StartCommand.class);
+  private static final WrappedLogger wrappedLog = new WrappedLogger(log);
   @Option(name = "--all", usage = "push all known projects")
   private boolean all;
 
   @Option(name = "--url", metaVar = "PATTERN", usage = "pattern to match URL on")
   private String urlMatch;
 
+  @Option(name = "--wait",
+      usage = "wait for replication to finish before exiting")
+  private boolean wait;
+
   @Argument(index = 0, multiValued = true, metaVar = "PROJECT", usage = "project name")
   private List<String> projectNames = new ArrayList<String>(2);
 
@@ -56,17 +66,62 @@ final class StartCommand extends SshCommand {
       throw new UnloggedFailure(1, "error: cannot combine --all and PROJECT");
     }
 
+    ReplicationState state =
+        new ReplicationState(ReplicationType.COMMAND, this);
+    Future<?> future = null;
     if (all) {
-      pushAllFactory.create(urlMatch).schedule(0, TimeUnit.SECONDS);
-
+      future = pushAllFactory.create(urlMatch, state).schedule(0, TimeUnit.SECONDS);
     } else {
       for (String name : projectNames) {
         Project.NameKey key = new Project.NameKey(name);
         if (projectCache.get(key) != null) {
-          replication.scheduleFullSync(key, urlMatch);
+          replication.scheduleFullSync(key, urlMatch, state);
         } else {
-          throw new UnloggedFailure(1, "error: '" + name + "': not a Gerrit project");
+          writeStdErrSync("error: '" + name + "': not a Gerrit project");
+        }
+      }
+      state.markAllPushTasksScheduled();
+    }
+
+    if (wait) {
+      if (future != null) {
+        try {
+          future.get();
+        } catch (InterruptedException e) {
+          wrappedLog.error("Thread was interrupted while waiting for PushAll operation to finish", e, state);
+          return;
+        } catch (ExecutionException e) {
+          wrappedLog.error("An exception was thrown in PushAll operation", e, state);
+          return;
+        }
+      }
+
+      if (state.hasPushTask()) {
+        try {
+          state.waitForReplication();
+        } catch (InterruptedException e) {
+          writeStdErrSync("We are interrupted while waiting replication to complete");
         }
+      } else {
+        writeStdOutSync("Nothing to replicate");
+      }
+    }
+  }
+
+  public void writeStdOutSync(final String message) {
+    if (wait) {
+      synchronized (stdout) {
+        stdout.println(message);
+        stdout.flush();
+      }
+    }
+  }
+
+  public void writeStdErrSync(final String message) {
+    if (wait) {
+      synchronized (stderr) {
+        stderr.println(message);
+        stderr.flush();
       }
     }
   }