Address pr comments

Author: tonyxiaowei

dbsync/common/src/main/java/com/google/edwmigration/dbsync/common/ChecksumGenerator.java

@@ -28,7 +28,7 @@ public ChecksumGenerator(@NonNegative int blockSize) {
     return blockSize;
   }
 
-  public void generate(Consumer<Checksum> out, ByteSource in) throws IOException {
+  public void generate(ChecksumConsumer<IOException> out, ByteSource in) throws IOException {
     // This deliberately throws if size is not Present.
     Optional<Long> dataSizeOptional = in.sizeIfKnown();
     if (!dataSizeOptional.isPresent()) {
@@ -76,4 +76,9 @@ public void generate(Consumer<Checksum> out, ByteSource in) throws IOException {
       }
     }
   }
+
+  public interface ChecksumConsumer<X extends Exception> {
+
+    void accept(Checksum checksum) throws X;
+  }
 }

dbsync/common/src/main/java/com/google/edwmigration/dbsync/common/InstructionGenerator.java

@@ -38,7 +38,7 @@ private static Instruction newLiteralInstruction(byte[] literalBuffer, int liter
         .build();
   }
 
-  public void generate(Consumer<? super Instruction> out, ByteSource in,
+  public void generate(InstructionConsumer<? extends IOException> out, ByteSource in,
       List<? extends Checksum> checksums) throws IOException {
     Int2ObjectMap<Collection<Checksum>> checksumMap = new Int2ObjectOpenHashMap<>(checksums.size());
     for (Checksum c : checksums) {
@@ -74,9 +74,9 @@ public void generate(Consumer<? super Instruction> out, ByteSource in,
           // And the wisdom to know the difference.
           MATCH:
           if (cc != null) {
-            HashCode strongHashCode = rollingChecksum.getStrongHashCode();
+            String strongHashCode = rollingChecksum.getStrongHashCode().toString();
             for (Checksum c : cc) {
-              if (c.getStrongChecksum().equals(strongHashCode.toString())) {
+              if (c.getStrongChecksum().equals(strongHashCode)) {
                 if (literalBufferLength > 0) {
                   if (DEBUG) {
                     LOG.debug("Emitting pre-match literal");
@@ -86,8 +86,8 @@ public void generate(Consumer<? super Instruction> out, ByteSource in,
                 }
                 Instruction insn = Instruction.newBuilder()
                     .setBlockLocation(BlockLocation.newBuilder()
-                            .setBlockOffset(c.getBlockOffset())
-                            .setBlockLength(c.getBlockLength()))
+                        .setBlockOffset(c.getBlockOffset())
+                        .setBlockLength(c.getBlockLength()))
                     .build();
                 out.accept(insn);
                 continue STREAM;
@@ -128,4 +128,13 @@ public void generate(Consumer<? super Instruction> out, ByteSource in,
       }
     }
   }
+
+  public interface InstructionConsumer<X extends Exception> {
+
+    /**
+     * Accepts an instruction from the generator
+     */
+    void accept(Instruction instruction) throws X;
+  }
+
 }

dbsync/common/src/main/java/com/google/edwmigration/dbsync/common/InstructionReceiver.java

@@ -25,18 +25,13 @@ public class InstructionReceiver implements Closeable {
 
   private final OutputStream out;
   private final ByteSource in;
-  private InputStream sourceStream;
-  private long currentSourceOffset = 0;
   @CheckForSigned
   private long copyStart = -1;
   private long copyLength = 0;
 
-  public InstructionReceiver(@WillCloseWhenClosed OutputStream out, ByteSource in)
-      throws IOException {
+  public InstructionReceiver(@WillCloseWhenClosed OutputStream out, ByteSource in) {
     this.out = Preconditions.checkNotNull(out, "Output was null.");
     this.in = Preconditions.checkNotNull(in, "Input was null.");
-
-    this.sourceStream = in.openStream();
   }
 
   public void receive(Instruction instruction) throws IOException {
@@ -68,64 +63,21 @@ public void close() throws IOException {
       flushCopy();
     } finally {
       out.close();
-      sourceStream.close();
     }
   }
 
   private void flushCopy() throws IOException {
     if (copyStart == -1) {
       return;
     }
-    LOG.info(String.format("Reuse bytes from %d for %d bytes", copyStart, copyLength));
-
-    // If the requested block is before our current position,
-    // we must re-open the stream from the beginning.
-    if (copyStart < currentSourceOffset) {
-      if (DEBUG) {
-        LOG.info(String.format("Target offset %d smaller than current offset %d, reopen stream",
-            copyStart, currentSourceOffset));
-      }
-      sourceStream.close();
-      sourceStream = in.openStream();
-      currentSourceOffset = 0;
+    if (DEBUG) {
+      LOG.info(String.format("Reuse bytes from %d for %d bytes", copyStart, copyLength));
     }
-    // Skip forward to the desired start.
-    long toSkip = copyStart - currentSourceOffset;
-    skip(toSkip);
-    copy(copyLength);
+    in.slice(copyStart, copyLength).copyTo(out);
 
     // Clear pending copy.
     copyStart = -1;
     copyLength = 0;
   }
 
-  private void copy(long copyLength) throws IOException {
-    // Now copy exactly copyLength bytes.
-    byte[] buffer = new byte[8192];
-    long remaining = copyLength;
-    while (remaining > 0) {
-      int bytesToRead = (int) Math.min(buffer.length, remaining);
-      int read = sourceStream.read(buffer, 0, bytesToRead);
-      if (read == -1) {
-        throw new EOFException("Unexpected end of stream while copying " + copyLength
-            + " bytes starting at offset " + copyStart);
-      }
-      out.write(buffer, 0, read);
-      remaining -= read;
-      currentSourceOffset += read;
-    }
-  }
-
-  private void skip(long length) throws IOException {
-    long remaining = length;
-    while (remaining > 0) {
-      long skipped = sourceStream.skip(remaining);
-      if (skipped < remaining) {
-        throw new EOFException("Unexpected end of stream while skipping.");
-      }
-      remaining -= skipped;
-      currentSourceOffset += skipped;
-    }
-  }
-
 }

dbsync/gcsync/src/main/java/com/google/edwmigration/dbsync/gcsync/GcsyncClient.java

@@ -13,6 +13,7 @@
 import com.google.cloud.storage.Blob;
 import com.google.common.base.Preconditions;
 import com.google.common.hash.Hashing;
+import com.google.common.io.ByteSink;
 import com.google.common.io.ByteSource;
 import com.google.edwmigration.dbsync.common.InstructionGenerator;
 import com.google.edwmigration.dbsync.proto.Checksum;
@@ -21,6 +22,7 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.io.OutputStreamWriter;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.nio.file.DirectoryStream;
@@ -148,18 +150,16 @@ private void uploadJar() throws URISyntaxException, IOException {
   }
 
   private void uploadFilesToRsyncList() throws URISyntaxException, IOException {
-    Path tempFile = Files.createTempFile("filesToRsyncTmp", ".txt");
-    try (BufferedWriter writer = Files.newBufferedWriter(tempFile)) {
+    ByteSink byteSink = gcsStorage.newByteSink(
+        new URI(tmpBucket).resolve(Constants.FILES_TO_RSYNC_FILE_NAME));
+
+    try (BufferedWriter writer = new BufferedWriter(
+        new OutputStreamWriter(byteSink.openBufferedStream()))) {
       for (Path file : filesToRsync) {
         writer.write(file.getFileName().toString());
         writer.newLine();
       }
     }
-
-    gcsStorage.uploadFile(tempFile,
-        new URI(tmpBucket).resolve(Constants.FILES_TO_RSYNC_FILE_NAME));
-
-    Files.deleteIfExists(tempFile);
   }
 
   private void executeMainOnCloudRun(String mainClassPath)
@@ -183,39 +183,30 @@ private void downloadCheckSumFiles() throws URISyntaxException, IOException {
       ByteSource byteSource = gcsStorage.newByteSource(
           new URI(tmpBucket).resolve(checksumFileName));
       byteSource.copyTo(
-          com.google.common.io.Files.asByteSink(getTemporaryCheckSumFileName(file).toFile()));
+          com.google.common.io.Files.asByteSink(Util.getTemporaryCheckSumFilePath(file).toFile()));
     }
   }
 
   private void sendRsyncInstructions() throws IOException, URISyntaxException {
     for (Path file : filesToRsync) {
-      Path tmpCheckSumFile = getTemporaryCheckSumFileName(file);
+      Path tmpCheckSumFile = Util.getTemporaryCheckSumFilePath(file);
       try (OutputStream instructionSink = gcsStorage.newByteSink(
               new URI(tmpBucket).resolve(Util.getInstructionFileName(file.getFileName().toString())))
           .openBufferedStream()) {
         try (InputStream inputStream = Files.newInputStream(tmpCheckSumFile)) {
           List<Checksum> checksums = getChecksumsFromFile(inputStream);
           ByteSource fileInput = com.google.common.io.Files.asByteSource(file.toFile());
 
-          instructionGenerator.generate(instruction -> {
-            try {
-              instruction.writeDelimitedTo(instructionSink);
-            } catch (IOException e) {
-              throw new RuntimeException(e);
-            }
-          }, fileInput, checksums);
+          instructionGenerator.generate(instruction ->
+                  instruction.writeDelimitedTo(instructionSink)
+              , fileInput, checksums);
         }
       }
 
       Files.deleteIfExists(tmpCheckSumFile);
     }
   }
 
-  private static Path getTemporaryCheckSumFileName(Path file) {
-    return file.resolveSibling(
-        String.format("%s.%s", file.getFileName(), Constants.CHECK_SUM_FILE_SUFFIX));
-  }
-
   private void reconStructFiles()
       throws IOException, ExecutionException, InterruptedException, URISyntaxException {
     executeMainOnCloudRun(Constants.RECONSTRUCT_FILE_MAIN);

dbsync/gcsync/src/main/java/com/google/edwmigration/dbsync/gcsync/GenerateCheckSumMain.java

@@ -41,13 +41,8 @@ public static void main(String[] args) throws Exception {
       ByteSink byteSink = gcsStorage.newByteSink(
           new URI(tmpBucket).resolve(Util.getCheckSumFileName(file)));
       try (OutputStream bufferedOutputStream = byteSink.openBufferedStream()) {
-        checksumGenerator.generate(checksum -> {
-          try {
-            checksum.writeDelimitedTo(bufferedOutputStream);
-          } catch (IOException e) {
-            throw new RuntimeException("Failed to generate checksum", e);
-          }
-        }, byteSource);
+        checksumGenerator.generate(checksum ->
+            checksum.writeDelimitedTo(bufferedOutputStream), byteSource);
         logger.log(Level.INFO,
             String.format("Finished generating check sum for: %s", file));
       }

dbsync/gcsync/src/main/java/com/google/edwmigration/dbsync/gcsync/Util.java

@@ -4,6 +4,7 @@
 import java.io.BufferedReader;
 import java.io.IOException;
 import java.io.InputStreamReader;
+import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.List;
 
@@ -22,6 +23,10 @@ public static List<String> getListOfFiles(ByteSource byteSource) throws IOExcept
     return files;
   }
 
+  public static Path getTemporaryCheckSumFilePath(Path file) {
+    return file.resolveSibling(getCheckSumFileName(file.getFileName().toString()));
+  }
+
   public static String getCheckSumFileName(String fileName) {
     return String.format("%s.%s", fileName, Constants.CHECK_SUM_FILE_SUFFIX);
   }

dbsync/gcsync/src/test/java/com/google/edwmigration/dbsync/gcsync/GcsyncClientTest.java

@@ -165,31 +165,6 @@ public void tearDown() {
     // No special cleanup required; TemporaryFolder is auto-cleaned.
   }
 
-  @Test
-  public void testSyncFiles_LargeFileMd5Mismatch_RunsJobAndDeletesJob()
-      throws Exception {
-    // Suppose the large file on GCS has an MD5 mismatch
-    Blob mockBlob = mockGcsStorage.getBlob(TARGET_BUCKET, "large.txt");
-    reset(mockBlob);
-    when(mockBlob.getMd5()).thenReturn("some-other-md5");
-
-    try {
-      clientUnderTest.syncFiles();
-    } catch (Exception e) {
-      // It's supposed to fail at instruction generation step
-    }
-
-    List<Path> rsyncFiles = getPrivateList(clientUnderTest, "filesToRsync");
-    assertTrue("large.txt should be in the rsync list", rsyncFiles.contains(largeFile.toPath()));
-
-    verify(mockJobsClient, times(1)).createJobAsync(any(CreateJobRequest.class));
-
-    verify(mockJobsClient.runJobOperationCallable(), times(1))
-        .futureCall(any(RunJobRequest.class));
-
-    verify(mockJobsClient, times(1)).deleteJobAsync(any(JobName.class));
-  }
-
   @Test
   public void testSyncFiles_LargeFileMd5Matches_NoCloudRun() throws Exception {
     clientUnderTest.syncFiles();