1 /*
2  *  Licensed to the Apache Software Foundation (ASF) under one or more
3  *  contributor license agreements.  See the NOTICE file distributed with
4  *  this work for additional information regarding copyright ownership.
5  *  The ASF licenses this file to You under the Apache License, Version 2.0
6  *  (the "License"); you may not use this file except in compliance with
7  *  the License.  You may obtain a copy of the License at
8  *
9  *      http://www.apache.org/licenses/LICENSE-2.0
10  *
11  *  Unless required by applicable law or agreed to in writing, software
12  *  distributed under the License is distributed on an "AS IS" BASIS,
13  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *  See the License for the specific language governing permissions and
15  *  limitations under the License.
16  */
17 package org.apache.tomcat.util.net;
18
19 import java.io.IOException;
20 import java.nio.ByteBuffer;
21 import java.util.ArrayList;
22 import java.util.Iterator;
23 import java.util.List;
24 import java.util.concurrent.LinkedBlockingDeque;
25
26 import org.apache.tomcat.util.buf.ByteBufferHolder;
27
28 /**
29  * Provides an expandable set of buffers for writes. Non-blocking writes can be
30  * of any size and may not be able to be written immediately or wholly contained
31  * in the buffer used to perform the writes to the next layer. This class
32  * provides a buffering capability to allow such writes to return immediately
33  * and also allows for the user provided buffers to be re-used / recycled as
34  * required.
35  */
36 public class WriteBuffer {
37
38     private final int bufferSize;
39
40     private final LinkedBlockingDeque<ByteBufferHolder> buffers = new LinkedBlockingDeque<>();
41
42     public WriteBuffer(int bufferSize) {
43         this.bufferSize = bufferSize;
44     }
45
46     void clear() {
47         buffers.clear();
48     }
49
50     void add(byte[] buf, int offset, int length) {
51         ByteBufferHolder holder = getByteBufferHolder(length);
52         holder.getBuf().put(buf, offset, length);
53     }
54
55
56     public void add(ByteBuffer from) {
57         ByteBufferHolder holder = getByteBufferHolder(from.remaining());
58         holder.getBuf().put(from);
59     }
60
61
62     private ByteBufferHolder getByteBufferHolder(int capacity) {
63         ByteBufferHolder holder = buffers.peekLast();
64         if (holder == null || holder.isFlipped() || holder.getBuf().remaining() < capacity) {
65             ByteBuffer buffer = ByteBuffer.allocate(Math.max(bufferSize, capacity));
66             holder = new ByteBufferHolder(buffer, false);
67             buffers.add(holder);
68         }
69         return holder;
70     }
71
72
73     public boolean isEmpty() {
74         return buffers.isEmpty();
75     }
76
77
78     /**
79      * Create an array of ByteBuffers from the current WriteBuffer, prefixing
80      * that array with the provided ByteBuffers.
81      *
82      * @param prefixes The additional ByteBuffers to add to the start of the
83      *                 array
84      *
85      * @return an array of ByteBuffers from the current WriteBuffer prefixed by
86      *         the provided ByteBuffers
87      */
88     ByteBuffer[] toArray(ByteBuffer... prefixes) {
89         List<ByteBuffer> result = new ArrayList<>();
90         for (ByteBuffer prefix : prefixes) {
91             if (prefix.hasRemaining()) {
92                 result.add(prefix);
93             }
94         }
95         for (ByteBufferHolder buffer : buffers) {
96             buffer.flip();
97             result.add(buffer.getBuf());
98         }
99         buffers.clear();
100         return result.toArray(new ByteBuffer[result.size()]);
101     }
102
103
104     boolean write(SocketWrapperBase<?> socketWrapper, boolean blocking) throws IOException {
105         Iterator<ByteBufferHolder> bufIter = buffers.iterator();
106         boolean dataLeft = false;
107         while (!dataLeft && bufIter.hasNext()) {
108             ByteBufferHolder buffer = bufIter.next();
109             buffer.flip();
110             if (blocking) {
111                 socketWrapper.writeBlocking(buffer.getBuf());
112             } else {
113                 socketWrapper.writeNonBlockingInternal(buffer.getBuf());
114             }
115             if (buffer.getBuf().remaining() == 0) {
116                 bufIter.remove();
117             } else {
118                 dataLeft = true;
119             }
120         }
121         return dataLeft;
122     }
123
124
125     public boolean write(Sink sink, boolean blocking) throws IOException {
126         Iterator<ByteBufferHolder> bufIter = buffers.iterator();
127         boolean dataLeft = false;
128         while (!dataLeft && bufIter.hasNext()) {
129             ByteBufferHolder buffer = bufIter.next();
130             buffer.flip();
131             dataLeft = sink.writeFromBuffer(buffer.getBuf(), blocking);
132             if (!dataLeft) {
133                 bufIter.remove();
134             }
135         }
136         return dataLeft;
137     }
138
139
140     /**
141      * Interface implemented by clients of the WriteBuffer to enable data to be
142      * written back out from the buffer.
143      */
144     public interface Sink {
145         boolean writeFromBuffer(ByteBuffer buffer, boolean block) throws IOException;
146     }
147 }
148