Home | History | Annotate | Download | only in internal 
      1 /*
      2  * Copyright (C) 2013 Square, Inc.
      3  *
      4  * Licensed under the Apache License, Version 2.0 (the "License");
      5  * you may not use this file except in compliance with the License.
      6  * You may obtain a copy of the License at
      7  *
      8  *      http://www.apache.org/licenses/LICENSE-2.0
      9  *
     10  * Unless required by applicable law or agreed to in writing, software
     11  * distributed under the License is distributed on an "AS IS" BASIS,
     12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     13  * See the License for the specific language governing permissions and
     14  * limitations under the License.
     15  */
     16 package com.squareup.okhttp.internal;
     17 
     18 import java.io.ByteArrayOutputStream;
     19 import java.io.IOException;
     20 import java.io.OutputStream;
     21 
     22 import static com.squareup.okhttp.internal.Util.checkOffsetAndCount;
     23 
     24 /**
     25  * An output stream wrapper that recovers from failures in the underlying stream
     26  * by replacing it with another stream. This class buffers a fixed amount of
     27  * data under the assumption that failures occur early in a stream's life.
     28  * If a failure occurs after the buffer has been exhausted, no recovery is
     29  * attempted.
     30  *
     31  * <p>Subclasses must override {@link #replacementStream} which will request a
     32  * replacement stream each time an {@link IOException} is encountered on the
     33  * current stream.
     34  */
     35 public abstract class FaultRecoveringOutputStream extends AbstractOutputStream {
     36   private final int maxReplayBufferLength;
     37 
     38   /** Bytes to transmit on the replacement stream, or null if no recovery is possible. */
     39   private ByteArrayOutputStream replayBuffer;
     40   private OutputStream out;
     41 
     42   /**
     43    * @param maxReplayBufferLength the maximum number of successfully written
     44    *     bytes to buffer so they can be replayed in the event of an error.
     45    *     Failure recoveries are not possible once this limit has been exceeded.
     46    */
     47   public FaultRecoveringOutputStream(int maxReplayBufferLength, OutputStream out) {
     48     if (maxReplayBufferLength < 0) throw new IllegalArgumentException();
     49     this.maxReplayBufferLength = maxReplayBufferLength;
     50     this.replayBuffer = new ByteArrayOutputStream(maxReplayBufferLength);
     51     this.out = out;
     52   }
     53 
     54   @Override public final void write(byte[] buffer, int offset, int count) throws IOException {
     55     if (closed) throw new IOException("stream closed");
     56     checkOffsetAndCount(buffer.length, offset, count);
     57 
     58     while (true) {
     59       try {
     60         out.write(buffer, offset, count);
     61 
     62         if (replayBuffer != null) {
     63           if (count + replayBuffer.size() > maxReplayBufferLength) {
     64             // Failure recovery is no longer possible once we overflow the replay buffer.
     65             replayBuffer = null;
     66           } else {
     67             // Remember the written bytes to the replay buffer.
     68             replayBuffer.write(buffer, offset, count);
     69           }
     70         }
     71         return;
     72       } catch (IOException e) {
     73         if (!recover(e)) throw e;
     74       }
     75     }
     76   }
     77 
     78   @Override public final void flush() throws IOException {
     79     if (closed) {
     80       return; // don't throw; this stream might have been closed on the caller's behalf
     81     }
     82     while (true) {
     83       try {
     84         out.flush();
     85         return;
     86       } catch (IOException e) {
     87         if (!recover(e)) throw e;
     88       }
     89     }
     90   }
     91 
     92   @Override public final void close() throws IOException {
     93     if (closed) {
     94       return;
     95     }
     96     while (true) {
     97       try {
     98         out.close();
     99         closed = true;
    100         return;
    101       } catch (IOException e) {
    102         if (!recover(e)) throw e;
    103       }
    104     }
    105   }
    106 
    107   /**
    108    * Attempt to replace {@code out} with another equivalent stream. Returns true
    109    * if a suitable replacement stream was found.
    110    */
    111   private boolean recover(IOException e) {
    112     if (replayBuffer == null) {
    113       return false; // Can't recover because we've dropped data that we would need to replay.
    114     }
    115 
    116     while (true) {
    117       OutputStream replacementStream = null;
    118       try {
    119         replacementStream = replacementStream(e);
    120         if (replacementStream == null) {
    121           return false;
    122         }
    123         replaceStream(replacementStream);
    124         return true;
    125       } catch (IOException replacementStreamFailure) {
    126         // The replacement was also broken. Loop to ask for another replacement.
    127         Util.closeQuietly(replacementStream);
    128         e = replacementStreamFailure;
    129       }
    130     }
    131   }
    132 
    133   /**
    134    * Returns true if errors in the underlying stream can currently be recovered.
    135    */
    136   public boolean isRecoverable() {
    137     return replayBuffer != null;
    138   }
    139 
    140   /**
    141    * Replaces the current output stream with {@code replacementStream}, writing
    142    * any replay bytes to it if they exist. The current output stream is closed.
    143    */
    144   public final void replaceStream(OutputStream replacementStream) throws IOException {
    145     if (!isRecoverable()) {
    146       throw new IllegalStateException();
    147     }
    148     if (this.out == replacementStream) {
    149       return; // Don't replace a stream with itself.
    150     }
    151     replayBuffer.writeTo(replacementStream);
    152     Util.closeQuietly(out);
    153     out = replacementStream;
    154   }
    155 
    156   /**
    157    * Returns a replacement output stream to recover from {@code e} thrown by the
    158    * previous stream. Returns a new OutputStream if recovery was successful, in
    159    * which case all previously-written data will be replayed. Returns null if
    160    * the failure cannot be recovered.
    161    */
    162   protected abstract OutputStream replacementStream(IOException e) throws IOException;
    163 }
    164