From 22e45daefd3a3cd13a5d2e5e1c4e3533582b82dc Mon Sep 17 00:00:00 2001
From: jwilson <jwilson@squareup.com>
Date: Wed, 17 Apr 2013 00:25:53 -0400
Subject: [PATCH] Recover from failed flushes and closes too.

---
 .../internal/FaultRecoveringOutputStream.java | 20 ++++++--
 .../FaultRecoveringOutputStreamTest.java      | 48 ++++++++++++++++++-
 2 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/okhttp/src/main/java/com/squareup/okhttp/internal/FaultRecoveringOutputStream.java b/okhttp/src/main/java/com/squareup/okhttp/internal/FaultRecoveringOutputStream.java
index d4e7507ab..0fbe7d618 100644
--- a/okhttp/src/main/java/com/squareup/okhttp/internal/FaultRecoveringOutputStream.java
+++ b/okhttp/src/main/java/com/squareup/okhttp/internal/FaultRecoveringOutputStream.java
@@ -84,15 +84,29 @@ public abstract class FaultRecoveringOutputStream extends OutputStream {
     if (closed) {
       return; // don't throw; this stream might have been closed on the caller's behalf
     }
-    out.flush();
+    while (true) {
+      try {
+        out.flush();
+        return;
+      } catch (IOException e) {
+        if (!recover(e)) throw e;
+      }
+    }
   }
 
   @Override public final void close() throws IOException {
     if (closed) {
       return;
     }
-    out.close();
-    closed = true;
+    while (true) {
+      try {
+        out.close();
+        closed = true;
+        return;
+      } catch (IOException e) {
+        if (!recover(e)) throw e;
+      }
+    }
   }
 
   /**
diff --git a/okhttp/src/test/java/com/squareup/okhttp/internal/FaultRecoveringOutputStreamTest.java b/okhttp/src/test/java/com/squareup/okhttp/internal/FaultRecoveringOutputStreamTest.java
index 5b3a621e4..e933c177c 100644
--- a/okhttp/src/test/java/com/squareup/okhttp/internal/FaultRecoveringOutputStreamTest.java
+++ b/okhttp/src/test/java/com/squareup/okhttp/internal/FaultRecoveringOutputStreamTest.java
@@ -47,7 +47,7 @@ public final class FaultRecoveringOutputStreamTest {
     }
   }
 
-  @Test public void successfulRecovery() throws Exception {
+  @Test public void successfulRecoveryOnWriteFault() throws Exception {
     FaultingOutputStream faulting1 = new FaultingOutputStream();
     FaultingOutputStream faulting2 = new FaultingOutputStream();
     TestFaultRecoveringOutputStream recovering = new TestFaultRecoveringOutputStream(10, faulting1);
@@ -70,6 +70,45 @@ public final class FaultRecoveringOutputStreamTest {
     assertEquals("abc", faulting2.receivedUtf8);
   }
 
+  @Test public void successfulRecoveryOnFlushFault() throws Exception {
+    FaultingOutputStream faulting1 = new FaultingOutputStream();
+    FaultingOutputStream faulting2 = new FaultingOutputStream();
+    TestFaultRecoveringOutputStream recovering = new TestFaultRecoveringOutputStream(10, faulting1);
+    recovering.replacements.addLast(faulting2);
+
+    recovering.write('a');
+    faulting1.nextFault = "bad weather";
+    recovering.flush();
+    assertEquals(Arrays.asList("bad weather"), recovering.exceptionMessages);
+    assertEquals("a", faulting1.receivedUtf8);
+    assertEquals("a", faulting2.receivedUtf8);
+    assertTrue(faulting1.closed);
+    assertFalse(faulting2.closed);
+    assertEquals("a", faulting2.flushedUtf8);
+
+    // Confirm that new data goes to the new stream.
+    recovering.write('b');
+    assertEquals("a", faulting1.receivedUtf8);
+    assertEquals("ab", faulting2.receivedUtf8);
+    assertEquals("a", faulting2.flushedUtf8);
+  }
+
+  @Test public void successfulRecoveryOnCloseFault() throws Exception {
+    FaultingOutputStream faulting1 = new FaultingOutputStream();
+    FaultingOutputStream faulting2 = new FaultingOutputStream();
+    TestFaultRecoveringOutputStream recovering = new TestFaultRecoveringOutputStream(10, faulting1);
+    recovering.replacements.addLast(faulting2);
+
+    recovering.write('a');
+    faulting1.nextFault = "termites";
+    recovering.close();
+    assertEquals(Arrays.asList("termites"), recovering.exceptionMessages);
+    assertEquals("a", faulting1.receivedUtf8);
+    assertEquals("a", faulting2.receivedUtf8);
+    assertTrue(faulting1.closed);
+    assertTrue(faulting2.closed);
+  }
+
   @Test public void replacementStreamFaultsImmediately() throws Exception {
     FaultingOutputStream faulting1 = new FaultingOutputStream();
     FaultingOutputStream faulting2 = new FaultingOutputStream();
@@ -145,6 +184,7 @@ public final class FaultRecoveringOutputStreamTest {
 
   static class FaultingOutputStream extends OutputStream {
     String receivedUtf8 = "";
+    String flushedUtf8 = null;
     String nextFault;
     boolean closed;
 
@@ -157,8 +197,14 @@ public final class FaultRecoveringOutputStreamTest {
       if (nextFault != null) throw new IOException(nextFault);
     }
 
+    @Override public void flush() throws IOException {
+      flushedUtf8 = receivedUtf8;
+      if (nextFault != null) throw new IOException(nextFault);
+    }
+
     @Override public void close() throws IOException {
       closed = true;
+      if (nextFault != null) throw new IOException(nextFault);
     }
   }