content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.hadoop.chukwa.datacollection.writer;
  
  import java.util.*;
  import java.util.concurrent.BlockingQueue;
  import java.util.concurrent.ArrayBlockingQueue;
  
  import org.apache.hadoop.chukwa.Chunk;
  import org.apache.hadoop.chukwa.util.Filter;
  import org.apache.hadoop.chukwa.util.RegexUtil.CheckedPatternSyntaxException;
  import org.apache.hadoop.conf.Configuration;
  import org.apache.log4j.Logger;
  
  import java.net.ServerSocket;
  import java.net.Socket;
  import java.nio.charset.Charset;
  import java.io.*;
  
  import org.apache.hadoop.chukwa.util.ExceptionUtil;
  
  /**
   * Effectively a "Tee" in the writer pipeline.
   * Accepts incoming connections on port specified by chukwaCollector.tee.port.
   * Defaults to 9094
   * 
   * Protocol is as follows:
   * Client ---&gt; TeeWriter   "RAW | WRITABLE &lt;filter&gt;" 
   *                  as per DumpChunks.
   *                  
   * TeeWriter ---&gt; Client "OK\n"                 
   *   In RAW mode               
   * TeeWriter ---&gt; Client (length(int)  byte[length])*
   *              An indefinite sequence of length, followed by byte array.
   *              
   *  In Writable mode
   * TeeWriter ---&gt; Client    (Chunk serialized as Writable)*
   *              An indefinite sequence of serialized chunks
   *              
   *  In English: clients should connect and say either "RAW " or "WRITABLE " 
   *  followed by a filter.  (Note that the keyword is followed by exactly one space.)
   *  They'll then receive either a sequence of byte arrays or of writable-serialized.
   *  
   *  Option chukwaCollector.tee.keepalive controls using TCP keepalive. Defaults to true.
   *  
   */
  public class SocketTeeWriter extends PipelineableWriter {
  
    public static final String WRITABLE = "WRITABLE";
    public static final String RAW = "RAW";
    public static final String ASCII_HEADER = "HEADER";
    
    static enum DataFormat {Raw, Writable, Header};
    
    static boolean USE_KEEPALIVE = true;
    public static final int DEFAULT_PORT = 9094;
    static int QUEUE_LENGTH = 1000;
    
    static Logger log = Logger.getLogger(SocketTeeWriter.class);
    volatile boolean running = true;
    int timeout;
  //  private final ExecutorService pool;
    
    /**
     * Listens for incoming connections, spawns a Tee to deal with each.
     */
    class SocketListenThread extends Thread {
      ServerSocket s;
      public SocketListenThread(Configuration conf) throws IOException {
        int portno = conf.getInt("chukwa.tee.port", DEFAULT_PORT);
        USE_KEEPALIVE = conf.getBoolean("chukwa.tee.keepalive", true);
        s = new ServerSocket(portno);
        setDaemon(true);
      }
      
      public void run() {
        log.info("listen thread started");
        try{
          while(running) {
            Socket sock = s.accept();
            log.info("got connection from " + sock.getInetAddress());
            new Tee(sock);
          }
        } catch(IOException e) {
          log.debug(ExceptionUtil.getStackTrace(e)); 
        }
      }
      
      public void shutdown() {
        try{
          //running was set to false by caller.
          s.close(); //to break out of run loop
          this.interrupt();
        } catch(IOException e) {
          log.debug(ExceptionUtil.getStackTrace(e)); 
        }
      }
    }
    
    /////////////////Internal class Tee//////////////////////
    /**
     * Manages a single socket connection
     */
    class Tee implements Runnable {
      Socket sock;
      BufferedReader in;
      DataOutputStream out;
      Filter rules;
      DataFormat fmt;
      final BlockingQueue<Chunk> sendQ;
      
      public Tee(Socket s) throws IOException {
        sock = s;
        //now initialize asynchronously
        sendQ = new ArrayBlockingQueue<Chunk>(QUEUE_LENGTH);
        
        Thread t = new Thread(this);
        t.setDaemon(true);
        t.start();
      }
      
      public void run() {
        setup();
        try {
          while(sock.isConnected()) {
            Chunk c = sendQ.take();
            
            if(fmt == DataFormat.Raw) {
              byte[] data = c.getData();
              out.writeInt(data.length);
              out.write(data);
            } else if(fmt == DataFormat.Writable)
              c.write(out);
            else {
              byte[] data = c.getData();
              byte[] header = (c.getSource()+ " " + c.getDataType() + " " + c.getStreamName()+ " "+  
                  c.getSeqID()+"\n").getBytes(Charset.forName("UTF-8")); 
              out.writeInt(data.length+ header.length);
              out.write(header);
              out.write(data);
            }
          }
          out.flush();
        } catch(IOException e) {
          log.info("lost tee: "+ e.toString());
          synchronized(tees) {
            tees.remove(this);
          }
        } catch(InterruptedException e) {
          //exit quietly
        }
      }
      
      /**
       * initializes the tee.
       */
      public void setup() {
        try {   //outer try catches IOExceptions
         try { //inner try catches bad command syntax errors
          sock.setSoTimeout(timeout);
          sock.setKeepAlive(USE_KEEPALIVE);
          in = new BufferedReader(new InputStreamReader(sock.getInputStream(), Charset.forName("UTF-8")));
          out = new DataOutputStream(sock.getOutputStream());
          String cmd = in.readLine();
          if(cmd==null) {
            throw new IllegalArgumentException("No input found.");
          }
          if(!cmd.contains(" ")) {
            
            throw new IllegalArgumentException(
                "command should be keyword pattern, but no ' ' seen: " + cmd);
          }
          String uppercased = cmd.substring(0, cmd.indexOf(' ')).toUpperCase();
          if(RAW.equals(uppercased))
            fmt = DataFormat.Raw;
          else if(WRITABLE.equals(uppercased))
            fmt = DataFormat.Writable;
          else if(ASCII_HEADER.equals(uppercased))
            fmt = DataFormat.Header;
          else {
            throw new IllegalArgumentException("bad command '" + uppercased+
                "' -- starts with neither '"+ RAW+ "' nor '"+ WRITABLE + " nor " 
                + ASCII_HEADER+"':" + cmd);
          }
          
          String cmdAfterSpace = cmd.substring(cmd.indexOf(' ')+1);
          if(cmdAfterSpace.toLowerCase().equals("all"))
            rules = Filter.ALL;
          else
            try {
              rules = new Filter(cmdAfterSpace);
            } catch (CheckedPatternSyntaxException pse) {
              out.write("Error parsing command as a regex: ".getBytes(Charset.forName("UTF-8")));
              out.write(pse.getMessage().getBytes(Charset.forName("UTF-8")));
              out.writeByte('\n');
              out.close();
              in.close();
              sock.close();
              log.warn(pse);
              return;
            }
  
            //now that we read everything OK we can add ourselves to list, and return.
          synchronized(tees) {
            tees.add(this);
          }
          out.write("OK\n".getBytes(Charset.forName("UTF-8")));
          log.info("tee to " + sock.getInetAddress() + " established");
        } catch(IllegalArgumentException e) {
            out.write(e.toString().getBytes(Charset.forName("UTF-8")));
            out.writeByte('\n');
            out.close();
            in.close();
            sock.close();
            log.warn(e);
          }//end inner catch
        } catch(IOException e) { //end outer catch
           log.warn(e);
        }
      }
      
      public void close() {
        try {
          out.close();
          in.close();
        } catch(Exception e) {
          log.debug(ExceptionUtil.getStackTrace(e));
        }
      }
  
      public void handle(Chunk c) {
        
        //don't ever block; just ignore this chunk if we don't have room for it.
        if(rules.matches(c)) {
          if(!sendQ.offer(c)) {
            log.debug("Queue is full.");
          }
        }
      }
    }
  
    /////////////////Main class SocketTeeWriter//////////////////////
    
    
    SocketListenThread listenThread;
    List<Tee> tees;
    
    @Override
    public void setNextStage(ChukwaWriter next) {
      this.next = next;
    }
  
    @Override
    public CommitStatus add(List<Chunk> chunks) throws WriterException {
      CommitStatus rv = ChukwaWriter.COMMIT_OK;
      if (next != null)
         rv = next.add(chunks); //pass data through
      synchronized(tees) {
        Iterator<Tee> loop = tees.iterator();
        while(loop.hasNext()) {
          Tee t = loop.next();
          for(Chunk c: chunks) {
            t.handle(c);
          }
        }
      }
      return rv;
    }
  
    @Override
    public void close() throws WriterException {
      if (next != null) 
  	next.close();
      running = false;
      listenThread.shutdown();
    }
  
    @Override
    public void init(Configuration c) throws WriterException {
      try {
        listenThread = new SocketListenThread(c);
        listenThread.start();
      } catch (IOException e) {
        throw new WriterException(e);
      }
      tees = new ArrayList<Tee>();
    }
  
  }
  