1 /*

     2  * Copyright (c) 2013, 2015, Oracle and/or its affiliates. All rights reserved.

     3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.

     4  *

     5  * This code is free software; you can redistribute it and/or modify it

     6  * under the terms of the GNU General Public License version 2 only, as

     7  * published by the Free Software Foundation.

     8  *

     9  * This code is distributed in the hope that it will be useful, but WITHOUT

    10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or

    11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License

    12  * version 2 for more details (a copy is included in the LICENSE file that

    13  * accompanied this code).

    14  *

    15  * You should have received a copy of the GNU General Public License version

    16  * 2 along with this work; if not, write to the Free Software Foundation,

    17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.

    18  *

    19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA

    20  * or visit www.oracle.com if you need additional information or have any

    21  * questions.

    22  */

    23 

    24 package jdk.test.lib.process;

    25 

    26 import java.io.BufferedInputStream;

    27 import java.io.ByteArrayOutputStream;

    28 import java.io.OutputStream;

    29 import java.io.InputStream;

    30 import java.io.IOException;

    31 import java.util.HashSet;

    32 import java.util.Set;

    33 import java.util.concurrent.Future;

    34 import java.util.concurrent.FutureTask;

    35 import java.util.concurrent.atomic.AtomicBoolean;

    36 

    37 public final class StreamPumper implements Runnable {

    38 

    39     private static final int BUF_SIZE = 256;

    40 

    41     /**

    42      * Pump will be called by the StreamPumper to process the incoming data

    43      */

    44     abstract public static class Pump {

    45         abstract void register(StreamPumper d);

    46     }

    47 

    48     /**

    49      * OutputStream -> Pump adapter

    50      */

    51     final public static class StreamPump extends Pump {

    52         private final OutputStream out;

    53         public StreamPump(OutputStream out) {

    54             this.out = out;

    55         }

    56 

    57         @Override

    58         void register(StreamPumper sp) {

    59             sp.addOutputStream(out);

    60         }

    61     }

    62 

    63     /**

    64      * Used to process the incoming data line-by-line

    65      */

    66     abstract public static class LinePump extends Pump {

    67         @Override

    68         final void register(StreamPumper sp) {

    69             sp.addLineProcessor(this);

    70         }

    71 

    72         abstract protected void processLine(String line);

    73     }

    74 

    75     private final InputStream in;

    76     private final Set<OutputStream> outStreams = new HashSet<>();

    77     private final Set<LinePump> linePumps = new HashSet<>();

    78 

    79     private final AtomicBoolean processing = new AtomicBoolean(false);

    80     private final FutureTask<Void> processingTask = new FutureTask<>(this, null);

    81 

    82     public StreamPumper(InputStream in) {

    83         this.in = in;

    84     }

    85 

    86     /**

    87      * Create a StreamPumper that reads from in and writes to out.

    88      *

    89      * @param in The stream to read from.

    90      * @param out The stream to write to.

    91      */

    92     public StreamPumper(InputStream in, OutputStream out) {

    93         this(in);

    94         this.addOutputStream(out);

    95     }

    96 

    97     /**

    98      * Implements Thread.run(). Continuously read from {@code in} and write to

    99      * {@code out} until {@code in} has reached end of stream. Abort on

   100      * interruption. Abort on IOExceptions.

   101      */

   102     @Override

   103     public void run() {

   104         try (BufferedInputStream is = new BufferedInputStream(in)) {

   105             ByteArrayOutputStream lineBos = new ByteArrayOutputStream();

   106             byte[] buf = new byte[BUF_SIZE];

   107             int len = 0;

   108             int linelen = 0;

   109 

   110             while ((len = is.read(buf)) > 0 && !Thread.interrupted()) {

   111                 for(OutputStream out : outStreams) {

   112                     out.write(buf, 0, len);

   113                 }

   114                 if (!linePumps.isEmpty()) {

   115                     int i = 0;

   116                     int lastcrlf = -1;

   117                     while (i < len) {

   118                         if (buf[i] == '\n' || buf[i] == '\r') {

   119                             int bufLinelen = i - lastcrlf - 1;

   120                             if (bufLinelen > 0) {

   121                                 lineBos.write(buf, lastcrlf + 1, bufLinelen);

   122                             }

   123                             linelen += bufLinelen;

   124 

   125                             if (linelen > 0) {

   126                                 lineBos.flush();

   127                                 final String line = lineBos.toString();

   128                                 linePumps.stream().forEach((lp) -> {

   129                                     lp.processLine(line);

   130                                 });

   131                                 lineBos.reset();

   132                                 linelen = 0;

   133                             }

   134                             lastcrlf = i;

   135                         }

   136 

   137                         i++;

   138                     }

   139                     if (lastcrlf == -1) {

   140                         lineBos.write(buf, 0, len);

   141                         linelen += len;

   142                     } else if (lastcrlf < len - 1) {

   143                         lineBos.write(buf, lastcrlf + 1, len - lastcrlf - 1);

   144                         linelen += len - lastcrlf - 1;

   145                     }

   146                 }

   147             }

   148 

   149         } catch (IOException e) {

   150             e.printStackTrace();

   151         } finally {

   152             for(OutputStream out : outStreams) {

   153                 try {

   154                     out.flush();

   155                 } catch (IOException e) {}

   156             }

   157             try {

   158                 in.close();

   159             } catch (IOException e) {}

   160         }

   161     }

   162 

   163     final void addOutputStream(OutputStream out) {

   164         outStreams.add(out);

   165     }

   166 

   167     final void addLineProcessor(LinePump lp) {

   168         linePumps.add(lp);

   169     }

   170 

   171     final public StreamPumper addPump(Pump ... pump) {

   172         if (processing.get()) {

   173             throw new IllegalStateException("Can not modify pumper while " +

   174                                             "processing is in progress");

   175         }

   176         for(Pump p : pump) {

   177             p.register(this);

   178         }

   179         return this;

   180     }

   181 

   182     final public Future<Void> process() {

   183         if (!processing.compareAndSet(false, true)) {

   184             throw new IllegalStateException("Can not re-run the processing");

   185         }

   186         Thread t = new Thread(new Runnable() {

   187             @Override

   188             public void run() {

   189                 processingTask.run();

   190             }

   191         });

   192         t.setDaemon(true);

   193         t.start();

   194 

   195         return processingTask;

   196     }

   197 }