How to avoid writing too much data to output and which blocks the thread?

Question

In a task which is run by a thread pool, I want to write a bunch of strings to remote, and there is a flag to indicate if the task has been cancelled or not.

I'm using the following code to make sure I can stop as soon as possible:

public void sendToRemote(Iterator<String> data, OutputStream remote) {
    try {
        System.out.println("##### staring sending")
        while(!this.cancelled && data.hasNext()) {
            remote.write(data.next())
        }
        System.out.println("##### finished sending")
        System.out.println("")
    } catch(Throwable e) {
        e.printStackTrace();
    } finally {
        remote.close();
    }
}

I found sometimes, if I give a very large data(or infinite iterator) to this method, even if I have set the this.cancelled to true later, it can't finish in time. The code seems blocked, and after a long time(1 minute or so), there will be an error like:

java.net.SocketTimeoutException: write blocked too long

So I guess it might be the remote.write method can block itself if there is too much data to send, but the remote doesn't consume it in time. Although I set this.cancelled to true , but the method has been blocked in the remote.write(data.next()) line for a long time, so it doesn't have chance to check the value of this.cancelled and skip the loop. Instead, at last, it throws an SocketTimeoutException after a long time.

Is my understanding correct? If it is, how can I avoid the blocking if there is too much data to send?

Answer 1

Try simply closing the remote OutputStream. Your thread will end with an exception.

1. Thread#1 busy performing sendToRemote();
2. Thread#2 decided enough is enough and closed the remote. (Assuming the OutPutStream object was not thread-local, as in a global reference somewhere)
3. Thread#1 died with an exception :)

EDIT I found this on the internet

Enabling linger and setting the timeout to a certain number of seconds will cause a subsequent call to Socket.Close to block until either all data in the send buffer has been sent or the timeout has elapsed.

Answer 2

The proper solution is probably to use NIO in some manner. I already commented on how Hadoop did it here using nio underneath .

But the simpler solution is in Dexter's answer . I also came across an answer from EJP who suggest to use a BufferedOutputStream to control when data goes out. So I combined the two to arrive at the TimedOutputStream shown below. It does not give complete control on output buffering to remote (much of it is done by the OS), but combining an appropriate buffer size and write timeout provides at least some control (see the second program for testing the TimedOutputStream ).

I have not completely tested the TimedOutputStream , so do your own due diligence.
updated write-method for better correlation between buffer-size and write timeout, also tweaked test program. 更新的write-method，用于缓冲区大小和写入超时之间的更好关联，还调整了测试程序。 Added comments about non-safe async close of socket outputstream.

import java.io.*;
import java.util.concurrent.*;

/**
 * A {@link BufferedOutputStream} that sets time-out tasks on write operations 
 * (typically when the buffer is flushed). If a write timeout occurs, the underlying outputstream is closed
 * (which may not be appropriate when sockets are used, see also comments on {@link TimedOutputStream#interruptWriteOut}).
 * A {@link ScheduledThreadPoolExecutor} is required to schedule the time-out tasks.
 * This {@code ScheduledThreadPoolExecutor} should have {@link ScheduledThreadPoolExecutor#setRemoveOnCancelPolicy(boolean)}
 * set to {@code true} to prevent a huge task queue.
 * If no {@code ScheduledThreadPoolExecutor} is provided in the constructor, 
 * the executor is created and shutdown with the {@link #close()} method. 
 * @author vanOekel
 *
 */
public class TimedOutputStream extends FilterOutputStream {

    protected int timeoutMs = 50_000;
    protected final boolean closeExecutor;
    protected final ScheduledExecutorService executor;
    protected ScheduledFuture<?> timeoutTask;
    protected volatile boolean writeTimedout;
    protected volatile IOException writeTimeoutCloseException;

    /* *** new methods not in BufferedOutputStream *** */

    /**
     * Default timeout is 50 seconds.
     */
    public void setTimeoutMs(int timeoutMs) {
        this.timeoutMs = timeoutMs;
    }

    public int getTimeoutMs() {
        return timeoutMs;
    }

    public boolean isWriteTimeout() {
        return writeTimedout;
    }

    /**
     * If a write timeout occurs and closing the underlying output-stream caused an exception,
     * then this method will return a non-null IOException.
     */
    public IOException getWriteTimeoutCloseException() {
        return writeTimeoutCloseException;
    }

    public ScheduledExecutorService getScheduledExecutor() {
        return executor;
    }

    /**
     * See {@link BufferedOutputStream#close()}.
     */
    @Override
    public void close() throws IOException {

        try {
            super.close(); // calls flush via FilterOutputStream.
        } finally {
            if (closeExecutor) {
                executor.shutdownNow();
            }
        }
    }

    /* ** Mostly a copy of java.io.BufferedOutputStream and updated with time-out options. *** */

    protected byte buf[];
    protected int count;

    public TimedOutputStream(OutputStream out) {
        this(out, null);
    }

    public TimedOutputStream(OutputStream out, ScheduledExecutorService executor) {
        this(out, 8192, executor);
    }

    public TimedOutputStream(OutputStream out, int size) {
        this(out, size, null);
    }

    public TimedOutputStream(OutputStream out, int size, ScheduledExecutorService executor) {
        super(out);
        if (size <= 0) {
            throw new IllegalArgumentException("Buffer size <= 0");
        }
        if (executor == null) {
            this.executor = Executors.newScheduledThreadPool(1);
            ScheduledThreadPoolExecutor stp = (ScheduledThreadPoolExecutor) this.executor;
            stp.setRemoveOnCancelPolicy(true);
            closeExecutor = true;
        } else {
            this.executor = executor;
            closeExecutor = false;
        }
        buf = new byte[size];
    }

    /**
     * Flushbuffer is called by all the write-methods and "flush()".
     */
    protected void flushBuffer(boolean flushOut) throws IOException {

        if (count > 0 || flushOut) {
            timeoutTask = executor.schedule(new TimeoutTask(this), getTimeoutMs(), TimeUnit.MILLISECONDS);
            try {
                // long start = System.currentTimeMillis(); int len = count;
                if (count > 0) {
                    out.write(buf, 0, count);
                    count = 0;
                }
                if (flushOut) {
                    out.flush(); // in case out is also buffered, this will do the actual write.
                }
                // System.out.println(Thread.currentThread().getName() + " Write [" + len + "] " + (flushOut ? "and flush " : "") + "time: " + (System.currentTimeMillis() - start));
            } finally {
                timeoutTask.cancel(false);
            }
        }
    }

    protected class TimeoutTask implements Runnable {

        protected final TimedOutputStream tout;
        public TimeoutTask(TimedOutputStream tout) {
            this.tout = tout;
        }

        @Override public void run() {
            tout.interruptWriteOut();
        }
    }

    /**
     * Closes the outputstream after a write timeout. 
     * If sockets are used, calling {@link java.net.Socket#shutdownOutput()} is probably safer
     * since the behavior of an async close of the outputstream is undefined. 
     */
    protected void interruptWriteOut() {

        try {
            writeTimedout = true;
            out.close();
        } catch (IOException e) {
            writeTimeoutCloseException = e;
        }
    }

    /**
     * See {@link BufferedOutputStream#write(int b)}
     */
    @Override
    public void write(int b) throws IOException {

        if (count >= buf.length) {
            flushBuffer(false);
        }
        buf[count++] = (byte)b;
    }

    /**
     * Like {@link BufferedOutputStream#write(byte[], int, int)}
     * but with one big difference: the full buffer is always written
     * to the underlying outputstream. Large byte-arrays are chopped
     * into buffer-size pieces and writtten out piece by piece.
     * <br>This provides a closer relation to the write timeout
     * and the maximum (buffer) size of the write-operation to wait on. 
     */
    @Override
    public void write(byte b[], int off, int len) throws IOException {

        if (count >= buf.length) {
            flushBuffer(false);
        }
        if (len <= buf.length - count) {
            System.arraycopy(b, off, buf, count, len);
            count += len;
        } else {
            final int fill = buf.length - count;
            System.arraycopy(b, off, buf, count, fill);
            count += fill;
            flushBuffer(false);
            final int remaining = len - fill;
            int start = off + fill;
            for (int i = 0; i < remaining / buf.length; i++) {
                System.arraycopy(b, start, buf, count, buf.length);
                count = buf.length;
                flushBuffer(false);
                start += buf.length;
            }
            count = remaining % buf.length;
            System.arraycopy(b, start, buf, 0, count);
        }
    }

    /**
     * See {@link BufferedOutputStream#flush()}
     * <br>If a write timeout occurred (i.e. {@link #isWriteTimeout()} returns {@code true}),
     * then this method does nothing. 
     */
    @Override
    public void flush() throws IOException {

        // Protect against flushing before closing after a write-timeout.
        // If that happens, then "out" is already closed in interruptWriteOut.
        if (!isWriteTimeout()) {
            flushBuffer(true);
        }
    }

}

And the test program:

import java.io.*;
import java.net.*;
import java.util.concurrent.*;

public class TestTimedSocketOut implements Runnable, Closeable {

    public static void main(String[] args) {

        TestTimedSocketOut m = new TestTimedSocketOut();
        try {
            m.run();
        } finally {
            m.close();
        }
    }

    final int clients = 3; // 2 is minimum, client 1 is expected to fail.
    final int timeOut = 1000;
    final int bufSize = 4096;
    final long maxWait = 5000L;
    // need a large array to write, else the OS just buffers everything and makes it work
    byte[] largeMsg = new byte[28_602];
    final ThreadPoolExecutor tp = (ThreadPoolExecutor) Executors.newCachedThreadPool();
    final ScheduledThreadPoolExecutor stp = (ScheduledThreadPoolExecutor) Executors.newScheduledThreadPool(1);
    final ConcurrentLinkedQueue<Closeable> closeables = new ConcurrentLinkedQueue<Closeable>();
    final CountDownLatch[] serversReady = new CountDownLatch[clients];
    final CountDownLatch clientsDone = new CountDownLatch(clients);
    final CountDownLatch serversDone = new CountDownLatch(clients);

    ServerSocket ss;
    int port;

    @Override public void run()  {

        stp.setRemoveOnCancelPolicy(true);
        try {
            ss = new ServerSocket();
            ss.bind(null);
            port = ss.getLocalPort();
            tp.execute(new SocketAccept());
            for (int i = 0; i < clients; i++) {
                serversReady[i] = new CountDownLatch(1);
                ClientSideSocket css = new ClientSideSocket(i);
                closeables.add(css);
                tp.execute(css);
                // need sleep to ensure client 0 connects first.
                Thread.sleep(50L);
            }
            if (!clientsDone.await(maxWait, TimeUnit.MILLISECONDS)) {
                println("CLIENTS DID NOT FINISH");
            } else {
                if (!serversDone.await(maxWait, TimeUnit.MILLISECONDS)) {
                    println("SERVERS DID NOT FINISH");
                } else {
                    println("Finished");
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    } 

    @Override public void close() {

        try { if (ss != null) ss.close(); } catch (Exception ignored) {}
        Closeable c = null;
        while ((c = closeables.poll()) != null) {
            try { c.close(); } catch (Exception ignored) {}
        }
        tp.shutdownNow();
        println("Scheduled tasks executed: " + stp.getTaskCount() + ", max. threads: " + stp.getLargestPoolSize());
        stp.shutdownNow();
    }

    class SocketAccept implements Runnable {

        @Override public void run() {
            try {
                for (int i = 0; i < clients; i++) {
                    SeverSideSocket sss = new SeverSideSocket(ss.accept(), i);
                    closeables.add(sss);
                    tp.execute(sss);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    class SeverSideSocket implements Runnable, Closeable {

        Socket s;
        int number, cnumber;
        boolean completed;

        public SeverSideSocket(Socket s, int number) {
            this.s = s;
            this.number = number;
            cnumber = -1;
        }

        @Override public void run() {

            String t = "nothing";
            try {
                DataInputStream in = new DataInputStream(s.getInputStream());
                DataOutputStream out = new DataOutputStream(s.getOutputStream());
                serversReady[number].countDown();
                Thread.sleep(timeOut);
                t = in.readUTF();
                in.readFully(new byte[largeMsg.length], 0, largeMsg.length);
                t += in.readUTF();
                out.writeByte(1);
                out.flush();
                cnumber = in.readInt();
                completed = true;
            } catch (Exception e) {
                println("server side " + number + " stopped after " + e);
                // e.printStackTrace();
            } finally {
                println("server side " + number + " received: " + t);
                if (completed && cnumber != number) {
                    println("server side " + number + " expected client number " + number + " but got " + cnumber);
                }
                close();
                serversDone.countDown();
            }
        }

        @Override public void close() {
            TestTimedSocketOut.close(s);
            s = null;
        }
    }

    class ClientSideSocket implements Runnable, Closeable {

        Socket s;
        int number;

        public ClientSideSocket(int number) {
            this.number = number;
        }

        @SuppressWarnings("resource")
        @Override public void run() {

            Byte b = -1;
            TimedOutputStream tout = null;
            try {
                s = new Socket();
                s.connect(new InetSocketAddress(port));
                DataInputStream in = new DataInputStream(s.getInputStream());
                tout = new TimedOutputStream(s.getOutputStream(), bufSize, stp);
                if (number == 1) {
                    // expect fail
                    tout.setTimeoutMs(timeOut / 2);
                } else {
                    // expect all OK
                    tout.setTimeoutMs(timeOut * 2);
                }
                DataOutputStream out = new DataOutputStream(tout);
                if (!serversReady[number].await(maxWait, TimeUnit.MILLISECONDS)) {
                    throw new RuntimeException("Server side for client side " + number + " not ready.");
                }
                out.writeUTF("client side " + number + " starting transfer");
                out.write(largeMsg);
                out.writeUTF(" - client side " + number + " completed transfer");
                out.flush();
                b = in.readByte();
                out.writeInt(number);
                out.flush();
            } catch (Exception e) {
                println("client side " + number + " stopped after " + e);
                // e.printStackTrace();
            } finally {
                println("client side " + number + " result: " + b);
                if (tout != null) {
                    if (tout.isWriteTimeout()) {
                        println("client side " + number + " had write timeout, close exception: " + tout.getWriteTimeoutCloseException());
                    } else {
                        println("client side " + number + " had no write timeout");
                    }
                }
                close();
                clientsDone.countDown();
            }
        }

        @Override public void close() {
            TestTimedSocketOut.close(s);
            s = null;
        }
    }

    private static void close(Socket s) {
        try { if (s != null) s.close(); } catch (Exception ignored) {}
    }

    private static final long START_TIME = System.currentTimeMillis(); 

    private static void println(String msg) {
        System.out.println((System.currentTimeMillis() - START_TIME) + "\t " + msg);
    }

}