private void driveClientIO() throws IOException, GearmanException {
   for (GearmanJobServerSession sess : sessionsMap.values()) {
     int interestOps = SelectionKey.OP_READ;
     if (sess.sessionHasDataToWrite()) {
       interestOps |= SelectionKey.OP_WRITE;
     }
     try {
       sess.getSelectionKey().interestOps(interestOps);
     } catch (IllegalStateException ise) {
       LOG.warn("Unable to drive IO for session " + sess + "," + " skipping.", ise);
       continue;
     }
   }
   ioAvailable.selectNow();
   Set<SelectionKey> keys = ioAvailable.selectedKeys();
   LOG.trace(
       "Driving IO for client "
           + this
           + ". "
           + keys.size()
           + " session(s) currently available for IO");
   Iterator<SelectionKey> iter = keys.iterator();
   while (iter.hasNext()) {
     SelectionKey key = iter.next();
     GearmanJobServerSession s = sessionsMap.get(key);
     s.driveSessionIO();
   }
 }
  public boolean addServer(GearmanJobServerConnection conn)
      throws IllegalArgumentException, IllegalStateException {

    if (conn == null) {
      throw new IllegalArgumentException("Connection can not be null");
    }
    // this is a sub-optimal way to look for dups, but addJobServer
    // ops should be infrequent enough that this should be a big penalty
    for (GearmanJobServerSession sess : sessionMap.values()) {
      if (sess.getConnection().equals(conn)) {
        return true;
      }
    }

    GearmanJobServerSession session = new GearmanJobServerSession(conn);
    if (ioAvailable == null) {
      try {
        ioAvailable = Selector.open();
      } catch (IOException ioe) {
        LOG.warn("Failed to connect to job server " + conn + ".", ioe);
        return false;
      }
    }
    try {
      session.initSession(ioAvailable, this);
    } catch (IOException ioe) {
      LOG.warn("Failed to initialize session with job server " + conn + ".", ioe);
      return false;
    }
    SelectionKey key = session.getSelectionKey();
    if (key == null) {
      String msg =
          "Session "
              + session
              + " has a null "
              + "selection key. Server will not be added to worker.";
      LOG.warn(msg);
      throw new IllegalStateException(msg);
    }
    sessionMap.put(key, session);

    GearmanPacket p =
        new GearmanPacketImpl(
            GearmanPacketMagic.REQ, GearmanPacketType.SET_CLIENT_ID, ByteUtils.toUTF8Bytes(id));
    session.submitTask(new GearmanTask(p));

    for (FunctionDefinition def : functionMap.values()) {
      p = generateCanDoPacket(def);
      session.submitTask(new GearmanTask(p)); // NOPMD
    }

    p = new GearmanPacketImpl(GearmanPacketMagic.REQ, getGrabJobPacketType(), new byte[0]);
    GearmanTask gsr = new GearmanTask(new GrabJobEventHandler(session), p);
    taskMap.put(session, gsr);
    session.submitTask(gsr);

    LOG.debug("Added server " + conn + " to worker " + this);
    return true;
  }
 private void shutDownSession(GearmanJobServerSession s) {
   if (s.isInitialized()) {
     SelectionKey k = s.getSelectionKey();
     if (k != null) {
       sessionsMap.remove(k);
       k.cancel();
     }
     s.closeSession();
   }
   sessionJobsMap.remove(s);
 }
  public boolean addJobServer(GearmanJobServerConnection newconn, boolean serverForwardsExceptions)
      throws IllegalArgumentException, IllegalStateException {

    // TODO remove this restriction
    if (!(newconn instanceof GearmanNIOJobServerConnection)) {
      throw new IllegalArgumentException(
          "Client currently only "
              + "supports "
              + GearmanNIOJobServerConnection.class.getName()
              + " connections.");
    }

    GearmanNIOJobServerConnection conn = (GearmanNIOJobServerConnection) newconn;

    if (!runState.equals(state.RUNNING)) {
      throw new RejectedExecutionException("Client has been shutdown");
    }

    GearmanJobServerSession session = new GearmanJobServerSession(conn);
    if (sessionsMap.values().contains(session)) {
      LOG.debug(
          "The server "
              + newconn
              + " was previously "
              + "added to the client. Ignoring add request.");
      return true;
    }

    try {
      if (ioAvailable == null) {
        ioAvailable = Selector.open();
      }

      session.initSession(ioAvailable, this);
      SelectionKey key = session.getSelectionKey();
      sessionsMap.put(key, session);
    } catch (IOException ioe) {
      LOG.warn("Failed to connect to job server " + newconn + ".", ioe);
      return false;
    }
    if (serverForwardsExceptions && !setForwardExceptions(session)) {
      return false;
    }
    sessionJobsMap.put(session, new HashMap<JobHandle, GearmanJobImpl>());

    LOG.info("Added connection " + conn + " to client " + this);
    return true;
  }
  private GearmanJobServerSession getSessionForTask() throws IOException {
    if (sessionsMap.values().isEmpty()) {
      throw new IOException("No servers registered with client");
    }

    ArrayList<GearmanJobServerSession> sessions = new ArrayList<GearmanJobServerSession>();
    sessions.addAll(sessionsMap.values());
    Random rand = new Random(System.currentTimeMillis());
    int s = rand.nextInt(sessions.size());
    GearmanJobServerSession session = sessions.get(s);
    IOException ioe = new IOException("No servers available for client");
    for (int i = 0; i < sessions.size(); i++) {
      if (!session.isInitialized()) {
        try {
          session.initSession(ioAvailable, this);
          sessionsMap.put(session.getSelectionKey(), session);
        } catch (Exception e) {
          if (sessions.size() > 1) {
            // try next one
            int prev = s;
            s = (s + 1) % sessions.size();
            session = sessions.get(s);
            LOG.warn(
                "Got exception attempting to retrieve "
                    + "session for task. Try next server at "
                    + "pos={}, previous pos={}",
                s,
                prev);
            continue;
          } else {
            break;
          }
        }
      }
      return session;
    }
    throw ioe;
  }
  public void work() {
    if (!state.equals(State.IDLE)) {
      throw new IllegalStateException(
          "Can not call work while worker " + "is running or shutting down");
    }

    state = State.RUNNING;
    // a map keeping track of sessions with connection errors
    // (to avoid printing an error about them in every reconnect attempt)
    Map<GearmanJobServerSession, Boolean> havingConnectionError =
        new HashMap<GearmanJobServerSession, Boolean>();

    while (isRunning()) {

      // look for sessions which have been disconnected and attempt to reconnect.
      for (Iterator<GearmanJobServerSession> iter = sessionMap.values().iterator();
          iter.hasNext(); ) {
        GearmanJobServerSession sess = iter.next();
        if (!sess.isInitialized()) {
          try {

            // reconnect, unregister old selection key and register new one
            SelectionKey oldKey = sess.isInitialized() ? sess.getSelectionKey() : null;
            sess.initSession(ioAvailable, this);
            if (oldKey != null) {
              iter.remove();
            }
            sessionMap.put(sess.getSelectionKey(), sess);

            // register all functions with the newly reconnected server
            for (FunctionDefinition d : functionMap.values()) {
              GearmanTask gsr = new GearmanTask(null, generateCanDoPacket(d));
              sess.submitTask(gsr);
            }

            GearmanPacket p =
                new GearmanPacketImpl(
                    GearmanPacketMagic.REQ,
                    GearmanPacketType.SET_CLIENT_ID,
                    ByteUtils.toUTF8Bytes(id));
            sess.submitTask(new GearmanTask(p));

            GearmanTask sessTask =
                new GearmanTask(
                    new GrabJobEventHandler(sess),
                    new GearmanPacketImpl(
                        GearmanPacketMagic.REQ, getGrabJobPacketType(), new byte[0]));
            sess.submitTask(sessTask);
            sess.driveSessionIO();

            // log reconnection message
            if (havingConnectionError.get(sess)) {
              LOG.info("Re-established connection to " + sess.getConnection().toString());
            }
            havingConnectionError.put(sess, false);
          } catch (IOException e) {
            if (!havingConnectionError.get(sess)) {
              LOG.warn("Error connecting to " + sess + ", will keep trying..");
            }

            havingConnectionError.put(sess, true);

            try {
              Thread.sleep(50);
            } catch (InterruptedException e1) {
            }
          }
        } else {
          havingConnectionError.put(sess, false);
        }
      }

      for (GearmanJobServerSession sess : sessionMap.values()) {
        // if still disconnected, skip
        if (!sess.isInitialized()) {
          continue;
        }
        int interestOps = SelectionKey.OP_READ;
        if (sess.sessionHasDataToWrite()) {
          interestOps |= SelectionKey.OP_WRITE;
        }
        sess.getSelectionKey().interestOps(interestOps);
      }
      try {
        ioAvailable.select(1);
      } catch (IOException io) {
        LOG.warn("Receieved IOException while" + " selecting for IO", io);
      }

      for (SelectionKey key : ioAvailable.selectedKeys()) {
        GearmanJobServerSession sess = sessionMap.get(key);
        if (sess == null) {
          LOG.warn("Worker does not have " + "session for key " + key);
          continue;
        }
        if (!sess.isInitialized()) {
          continue;
        }
        try {
          GearmanTask sessTask = taskMap.get(sess);
          if (sessTask == null) {
            sessTask =
                new GearmanTask( // NOPMD
                    new GrabJobEventHandler(sess),
                    new GearmanPacketImpl(
                        GearmanPacketMagic.REQ, getGrabJobPacketType(), new byte[0]));
            taskMap.put(sess, sessTask);
            sess.submitTask(sessTask);
            LOG.debug(
                "Worker: "
                    + this
                    + " submitted a "
                    + sessTask.getRequestPacket().getPacketType()
                    + " to session: "
                    + sess);
          }
          sess.driveSessionIO();
          // For the time being we will execute the jobs synchronously
          // in the future, I expect to change this.
          if (!functionList.isEmpty()) {
            GearmanFunction fun = functionList.remove();
            submitFunction(fun);
            statistics();
          }

        } catch (IOException ioe) {
          LOG.warn("Received IOException while driving" + " IO on session " + sess, ioe);
          sess.closeSession();
          continue;
        }
      }
    }

    shutDownWorker(true);
  }