From 94dba1ede9b58fe025ca9a59371bcc24b1e1bdb3 Mon Sep 17 00:00:00 2001
From: Jesse Wilson <jwilson@squareup.com>
Date: Wed, 11 Dec 2019 09:11:20 -0500
Subject: [PATCH] Remove empty tail segment (1.x branch)

Manual backport of cketti's change
https://github.com/square/okio/pull/686/commits
---
 okio/src/main/java/okio/Buffer.java     |  5 +++++
 okio/src/main/java/okio/Okio.java       |  9 ++++++++-
 okio/src/test/java/okio/BufferTest.java |  7 +++++++
 okio/src/test/java/okio/OkioTest.java   | 12 ++++++++++++
 okio/src/test/java/okio/TestUtil.java   |  7 +++++++
 5 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/okio/src/main/java/okio/Buffer.java b/okio/src/main/java/okio/Buffer.java
index c168bc77f6..ea9aeffc0c 100644
--- a/okio/src/main/java/okio/Buffer.java
+++ b/okio/src/main/java/okio/Buffer.java
@@ -257,6 +257,11 @@ private void readFrom(InputStream in, long byteCount, boolean forever) throws IO
       int maxToCopy = (int) Math.min(byteCount, Segment.SIZE - tail.limit);
       int bytesRead = in.read(tail.data, tail.limit, maxToCopy);
       if (bytesRead == -1) {
+        if (tail.pos == tail.limit) {
+          // We allocated a tail segment, but didn't end up needing it. Recycle!
+          head = tail.pop();
+          SegmentPool.recycle(tail);
+        }
         if (forever) return;
         throw new EOFException();
       }
diff --git a/okio/src/main/java/okio/Okio.java b/okio/src/main/java/okio/Okio.java
index 1163ffada8..da56d876f8 100644
--- a/okio/src/main/java/okio/Okio.java
+++ b/okio/src/main/java/okio/Okio.java
@@ -138,7 +138,14 @@ private static Source source(final InputStream in, final Timeout timeout) {
           Segment tail = sink.writableSegment(1);
           int maxToCopy = (int) Math.min(byteCount, Segment.SIZE - tail.limit);
           int bytesRead = in.read(tail.data, tail.limit, maxToCopy);
-          if (bytesRead == -1) return -1;
+          if (bytesRead == -1) {
+            if (tail.pos == tail.limit) {
+              // We allocated a tail segment, but didn't end up needing it. Recycle!
+              sink.head = tail.pop();
+              SegmentPool.recycle(tail);
+            }
+            return -1;
+          }
           tail.limit += bytesRead;
           sink.size += bytesRead;
           return bytesRead;
diff --git a/okio/src/test/java/okio/BufferTest.java b/okio/src/test/java/okio/BufferTest.java
index c70623376d..dfaba64024 100644
--- a/okio/src/test/java/okio/BufferTest.java
+++ b/okio/src/test/java/okio/BufferTest.java
@@ -25,6 +25,7 @@
 import org.junit.Test;
 
 import static java.util.Arrays.asList;
+import static okio.TestUtil.assertNoEmptySegments;
 import static okio.TestUtil.bufferWithRandomSegmentLayout;
 import static okio.TestUtil.repeat;
 import static okio.Util.UTF_8;
@@ -292,6 +293,12 @@ private List<Integer> moveBytesBetweenBuffers(String... contents) throws IOExcep
     assertEquals("hello, wor", out);
   }
 
+  @Test public void readFromDoesNotLeaveEmptyTailSegment() throws IOException {
+    Buffer buffer = new Buffer();
+    buffer.readFrom(new ByteArrayInputStream(new byte[Segment.SIZE]));
+    assertNoEmptySegments(buffer);
+  }
+
   @Test public void moveAllRequestedBytesWithRead() throws Exception {
     Buffer sink = new Buffer();
     sink.writeUtf8(repeat('a', 10));
diff --git a/okio/src/test/java/okio/OkioTest.java b/okio/src/test/java/okio/OkioTest.java
index 73e9b35327..f47ba0e9ab 100644
--- a/okio/src/test/java/okio/OkioTest.java
+++ b/okio/src/test/java/okio/OkioTest.java
@@ -25,6 +25,7 @@
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
 
+import static okio.TestUtil.assertNoEmptySegments;
 import static okio.TestUtil.repeat;
 import static okio.Util.UTF_8;
 import static org.junit.Assert.assertEquals;
@@ -119,6 +120,17 @@ public final class OkioTest {
     assertEquals(-1, source.read(sink, 1));
   }
 
+  @Test public void sourceFromInputStreamWithSegmentSize() throws Exception {
+    InputStream in = new ByteArrayInputStream(new byte[Segment.SIZE]);
+    Source source = Okio.source(in);
+    Buffer sink = new Buffer();
+
+    assertEquals(Segment.SIZE, source.read(sink, Segment.SIZE));
+    assertEquals(-1, source.read(sink, Segment.SIZE));
+
+    assertNoEmptySegments(sink);
+  }
+
   @Test public void sourceFromInputStreamBounds() throws Exception {
     Source source = Okio.source(new ByteArrayInputStream(new byte[100]));
     try {
diff --git a/okio/src/test/java/okio/TestUtil.java b/okio/src/test/java/okio/TestUtil.java
index 0f355ffac3..c881a21c3f 100644
--- a/okio/src/test/java/okio/TestUtil.java
+++ b/okio/src/test/java/okio/TestUtil.java
@@ -24,12 +24,19 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 
 final class TestUtil {
   private TestUtil() {
   }
 
+  static void assertNoEmptySegments(Buffer buffer) {
+    for (Integer segmentSize : buffer.segmentSizes()) {
+      assertNotEquals("Expected all segments to be non-empty", 0, segmentSize.intValue());
+    }
+  }
+
   static void assertByteArraysEquals(byte[] a, byte[] b) {
     assertEquals(Arrays.toString(a), Arrays.toString(b));
   }