[Cronet] Add isDone() check in CronetBidirectionalStream#onWritevCompleted

components/cronet/android/test/javatests/src/org/chromium/net/TestBidirectionalStreamCallback.java

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 package org.chromium.net;
 
 import android.os.ConditionVariable;
 
 import static junit.framework.Assert.assertEquals;
 import static junit.framework.Assert.assertFalse;
@@ skipping 51 matching lines @@
 
     // position() of ByteBuffer prior to read() call.
     private int mBufferPositionBeforeRead;
 
     // Data to write.
     private final ArrayList<WriteBuffer> mWriteBuffers = new ArrayList<WriteBuffer>();
 
     // Buffers that we yet to receive the corresponding onWriteCompleted callback.
     private final ArrayList<WriteBuffer> mWriteBuffersToBeAcked = new ArrayList<WriteBuffer>();
 
+    // Whether to use a direct executor.
+    private final boolean mUseDirectExecutor;
+    private final DirectExecutor mDirectExecutor;
+
     private class ExecutorThreadFactory implements ThreadFactory {
         public Thread newThread(Runnable r) {
             mExecutorThread = new Thread(r);
             return mExecutorThread;
         }
     }
 
     private static class WriteBuffer {
         final ByteBuffer mBuffer;
         final boolean mFlush;
         public WriteBuffer(ByteBuffer buffer, boolean flush) {
             mBuffer = buffer;
             mFlush = flush;
         }
     }
 
+    private static class DirectExecutor implements Executor {
+        @Override
+        public void execute(Runnable task) {
+            task.run();
+        }
+    }
+
     public enum ResponseStep {
         NOTHING,
         ON_STREAM_READY,
         ON_RESPONSE_STARTED,
         ON_READ_COMPLETED,
         ON_WRITE_COMPLETED,
         ON_TRAILERS,
         ON_CANCELED,
         ON_FAILED,
         ON_SUCCEEDED,
     }
 
     public enum FailureType {
         NONE,
         CANCEL_SYNC,
         CANCEL_ASYNC,
         // Same as above, but continues to advance the stream after posting
         // the cancellation task.
         CANCEL_ASYNC_WITHOUT_PAUSE,
         THROW_SYNC
     }
 
+    public TestBidirectionalStreamCallback() {
+        mUseDirectExecutor = false;
+        mDirectExecutor = null;
+    }
+
+    public TestBidirectionalStreamCallback(boolean useDirectExecutor) {
+        mUseDirectExecutor = useDirectExecutor;
+        mDirectExecutor = new DirectExecutor();
+    }
+
     public void setAutoAdvance(boolean autoAdvance) {
         mAutoAdvance = autoAdvance;
     }
 
     public void setFailure(FailureType failureType, ResponseStep failureStep) {
         mFailureStep = failureStep;
         mFailureType = failureType;
     }
 
     public void blockForDone() {
         mDone.block();
     }
 
     public void waitForNextReadStep() {
         mReadStepBlock.block();
         mReadStepBlock.close();
     }
 
     public void waitForNextWriteStep() {
         mWriteStepBlock.block();
         mWriteStepBlock.close();
     }
 
     public Executor getExecutor() {
+        if (mUseDirectExecutor) {
+            return mDirectExecutor;
+        }
         return mExecutorService;
     }
 
     public void shutdownExecutor() {
+        if (mUseDirectExecutor) {
+            throw new UnsupportedOperationException("DirectExecutor doesn't support shutdown");
+        }
         mExecutorService.shutdown();
     }
 
     public void addWriteData(byte[] data) {
         addWriteData(data, true);
     }
 
     public void addWriteData(byte[] data, boolean flush) {
         ByteBuffer writeBuffer = ByteBuffer.allocateDirect(data.length);
         writeBuffer.put(data);
         writeBuffer.flip();
         mWriteBuffers.add(new WriteBuffer(writeBuffer, flush));
         mWriteBuffersToBeAcked.add(new WriteBuffer(writeBuffer, flush));
     }
 
     @Override
     public void onStreamReady(BidirectionalStream stream) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertFalse(stream.isDone());
         assertEquals(ResponseStep.NOTHING, mResponseStep);
         assertNull(mError);
         mResponseStep = ResponseStep.ON_STREAM_READY;
         if (maybeThrowCancelOrPause(stream, mWriteStepBlock)) {
             return;
         }
         startNextWrite(stream);
     }
 
     @Override
     public void onResponseHeadersReceived(BidirectionalStream stream, UrlResponseInfo info) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertFalse(stream.isDone());
         assertTrue(mResponseStep == ResponseStep.NOTHING
                 || mResponseStep == ResponseStep.ON_STREAM_READY
                 || mResponseStep == ResponseStep.ON_WRITE_COMPLETED);
         assertNull(mError);
 
         mResponseStep = ResponseStep.ON_RESPONSE_STARTED;
         mResponseInfo = info;
         if (maybeThrowCancelOrPause(stream, mReadStepBlock)) {
             return;
         }
         startNextRead(stream);
     }
 
     @Override
     public void onReadCompleted(BidirectionalStream stream, UrlResponseInfo info,
             ByteBuffer byteBuffer, boolean endOfStream) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertFalse(stream.isDone());
         assertTrue(mResponseStep == ResponseStep.ON_RESPONSE_STARTED
                 || mResponseStep == ResponseStep.ON_READ_COMPLETED
                 || mResponseStep == ResponseStep.ON_WRITE_COMPLETED
                 || mResponseStep == ResponseStep.ON_TRAILERS);
         assertNull(mError);
 
         mResponseStep = ResponseStep.ON_READ_COMPLETED;
         mResponseInfo = info;
 
@@ skipping 13 matching lines @@
         }
         // Do not read if EOF has been reached.
         if (!endOfStream) {
             startNextRead(stream);
         }
     }
 
     @Override
     public void onWriteCompleted(BidirectionalStream stream, UrlResponseInfo info,
             ByteBuffer buffer, boolean endOfStream) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertFalse(stream.isDone());
         assertNull(mError);
         mResponseStep = ResponseStep.ON_WRITE_COMPLETED;
         mResponseInfo = info;
         if (!mWriteBuffersToBeAcked.isEmpty()) {
             assertEquals(buffer, mWriteBuffersToBeAcked.get(0).mBuffer);
             mWriteBuffersToBeAcked.remove(0);
         }
         if (maybeThrowCancelOrPause(stream, mWriteStepBlock)) {
             return;
         }
         startNextWrite(stream);
     }
 
     @Override
     public void onResponseTrailersReceived(BidirectionalStream stream, UrlResponseInfo info,
             UrlResponseInfo.HeaderBlock trailers) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertFalse(stream.isDone());
         assertNull(mError);
         mResponseStep = ResponseStep.ON_TRAILERS;
         mResponseInfo = info;
         mTrailers = trailers;
         if (maybeThrowCancelOrPause(stream, mReadStepBlock)) {
             return;
         }
     }
 
     @Override
     public void onSucceeded(BidirectionalStream stream, UrlResponseInfo info) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertTrue(stream.isDone());
         assertTrue(mResponseStep == ResponseStep.ON_RESPONSE_STARTED
                 || mResponseStep == ResponseStep.ON_READ_COMPLETED
                 || mResponseStep == ResponseStep.ON_WRITE_COMPLETED
                 || mResponseStep == ResponseStep.ON_TRAILERS);
         assertFalse(mOnErrorCalled);
         assertFalse(mOnCanceledCalled);
         assertNull(mError);
         assertEquals(0, mWriteBuffers.size());
         assertEquals(0, mWriteBuffersToBeAcked.size());
 
         mResponseStep = ResponseStep.ON_SUCCEEDED;
         mResponseInfo = info;
         openDone();
         maybeThrowCancelOrPause(stream, mReadStepBlock);
     }
 
     @Override
     public void onFailed(BidirectionalStream stream, UrlResponseInfo info, CronetException error) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertTrue(stream.isDone());
         // Shouldn't happen after success.
         assertTrue(mResponseStep != ResponseStep.ON_SUCCEEDED);
         // Should happen at most once for a single stream.
         assertFalse(mOnErrorCalled);
         assertFalse(mOnCanceledCalled);
         assertNull(mError);
         mResponseStep = ResponseStep.ON_FAILED;
         mResponseInfo = info;
 
         mOnErrorCalled = true;
         mError = error;
         openDone();
         maybeThrowCancelOrPause(stream, mReadStepBlock);
     }
 
     @Override
     public void onCanceled(BidirectionalStream stream, UrlResponseInfo info) {
-        assertEquals(mExecutorThread, Thread.currentThread());
+        checkOnValidThread();
         assertTrue(stream.isDone());
         // Should happen at most once for a single stream.
         assertFalse(mOnCanceledCalled);
         assertFalse(mOnErrorCalled);
         assertNull(mError);
         mResponseStep = ResponseStep.ON_CANCELED;
         mResponseInfo = info;
 
         mOnCanceledCalled = true;
         openDone();
@@ skipping 65 matching lines @@
             }
         };
         if (mFailureType == FailureType.CANCEL_ASYNC
                 || mFailureType == FailureType.CANCEL_ASYNC_WITHOUT_PAUSE) {
             getExecutor().execute(task);
         } else {
             task.run();
         }
         return mFailureType != FailureType.CANCEL_ASYNC_WITHOUT_PAUSE;
     }
+
+    /**
+     * Checks whether callback methods are invoked on the correct thread.
+     */
+    private void checkOnValidThread() {
+        if (!mUseDirectExecutor) {
+            assertEquals(mExecutorThread, Thread.currentThread());
+        }
+    }
 }