1 /*
   2  * Copyright (c) 2013, 2016, 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     public abstract static class Pump {
  45         abstract void register(StreamPumper d);
  46     }
  47 
  48     /**
  49      * OutputStream -> Pump adapter
  50      */
  51     public final 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     public abstract static class LinePump extends Pump {
  67         @Override
  68         final void register(StreamPumper sp) {
  69             sp.addLineProcessor(this);
  70         }
  71 
  72         protected abstract 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 
  81     public StreamPumper(InputStream in) {
  82         this.in = in;
  83     }
  84 
  85     /**
  86      * Create a StreamPumper that reads from in and writes to out.
  87      *
  88      * @param in The stream to read from.
  89      * @param out The stream to write to.
  90      */
  91     public StreamPumper(InputStream in, OutputStream out) {
  92         this(in);
  93         this.addOutputStream(out);
  94     }
  95 
  96     /**
  97      * Implements Thread.run(). Continuously read from {@code in} and write to
  98      * {@code out} until {@code in} has reached end of stream. Abort on
  99      * interruption. Abort on IOExceptions.
 100      */
 101     @Override
 102     public void run() {
 103         try (BufferedInputStream is = new BufferedInputStream(in)) {
 104             ByteArrayOutputStream lineBos = new ByteArrayOutputStream();
 105             byte[] buf = new byte[BUF_SIZE];
 106             int len = 0;
 107             int linelen = 0;
 108 
 109             while ((len = is.read(buf)) > 0 && !Thread.interrupted()) {
 110                 for (OutputStream out : outStreams) {
 111                     out.write(buf, 0, len);
 112                 }
 113                 if (!linePumps.isEmpty()) {
 114                     int i = 0;
 115                     int lastcrlf = -1;
 116                     while (i < len) {
 117                         if (buf[i] == '\n' || buf[i] == '\r') {
 118                             int bufLinelen = i - lastcrlf - 1;
 119                             if (bufLinelen > 0) {
 120                                 lineBos.write(buf, lastcrlf + 1, bufLinelen);
 121                             }
 122                             linelen += bufLinelen;
 123 
 124                             if (linelen > 0) {
 125                                 lineBos.flush();
 126                                 final String line = lineBos.toString();
 127                                 linePumps.forEach((lp) -> lp.processLine(line));
 128                                 lineBos.reset();
 129                                 linelen = 0;
 130                             }
 131                             lastcrlf = i;
 132                         }
 133 
 134                         i++;
 135                     }
 136                     if (lastcrlf == -1) {
 137                         lineBos.write(buf, 0, len);
 138                         linelen += len;
 139                     } else if (lastcrlf < len - 1) {
 140                         lineBos.write(buf, lastcrlf + 1, len - lastcrlf - 1);
 141                         linelen += len - lastcrlf - 1;
 142                     }
 143                 }
 144             }
 145 
 146         } catch (IOException e) {
 147             e.printStackTrace();
 148         } finally {
 149             for (OutputStream out : outStreams) {
 150                 try {
 151                     out.flush();
 152                 } catch (IOException e) {}
 153             }
 154             try {
 155                 in.close();
 156             } catch (IOException e) {}
 157         }
 158     }
 159 
 160     final void addOutputStream(OutputStream out) {
 161         outStreams.add(out);
 162     }
 163 
 164     final void addLineProcessor(LinePump lp) {
 165         linePumps.add(lp);
 166     }
 167 
 168     public final StreamPumper addPump(Pump ... pump) {
 169         if (processing.get()) {
 170             throw new IllegalStateException("Can not modify pumper while " +
 171                                             "processing is in progress");
 172         }
 173         for (Pump p : pump) {
 174             p.register(this);
 175         }
 176         return this;
 177     }
 178 
 179     public final Future<Void> process() {
 180         if (!processing.compareAndSet(false, true)) {
 181             throw new IllegalStateException("Can not re-run the processing");
 182         }
 183         FutureTask<Void> result = new FutureTask<>(this, null);
 184         Thread t = new Thread(result);
 185         t.setDaemon(true);
 186         t.start();
 187 
 188         return result;
 189     }
 190 }