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 addNewJob(GearmanSessionEvent event) {
   byte[] handle, data, functionNameBytes, unique;
   GearmanPacket p = event.getPacket();
   GearmanJobServerSession sess = event.getSession();
   String functionName;
   handle = p.getDataComponentValue(GearmanPacket.DataComponentName.JOB_HANDLE);
   functionNameBytes = p.getDataComponentValue(GearmanPacket.DataComponentName.FUNCTION_NAME);
   data = p.getDataComponentValue(GearmanPacket.DataComponentName.DATA);
   unique = p.getDataComponentValue(DataComponentName.UNIQUE_ID);
   functionName = ByteUtils.fromUTF8Bytes(functionNameBytes);
   FunctionDefinition def = functionMap.get(functionName);
   if (def == null) {
     GearmanTask gsr =
         new GearmanTask(
             new GearmanPacketImpl(GearmanPacketMagic.REQ, GearmanPacketType.WORK_FAIL, handle));
     sess.submitTask(gsr);
   } else {
     GearmanFunction function = def.getFactory().getFunction();
     function.setData(data);
     function.setJobHandle(handle);
     function.registerEventListener(sess);
     if (unique != null && unique.length > 0) {
       function.setUniqueId(unique);
     }
     functionList.add(function);
   }
 }
 public void handleSessionEvent(GearmanSessionEvent event)
     throws IllegalArgumentException, IllegalStateException {
   GearmanPacket p = event.getPacket();
   GearmanJobServerSession s = event.getSession();
   GearmanPacketType t = p.getPacketType();
   LOG.debug(
       "Worker "
           + this
           + " handling session event"
           + " ( Session = "
           + s
           + " Event = "
           + t
           + " )");
   switch (t) {
     case JOB_ASSIGN:
       // TODO Figure out what the right behavior is if JobUUIDRequired was false when we submitted
       // but is now true
       taskMap.remove(s);
       addNewJob(event);
       break;
     case JOB_ASSIGN_UNIQ:
       // TODO Figure out what the right behavior is if JobUUIDRequired was true when we submitted
       // but is now false
       taskMap.remove(s);
       addNewJob(event);
       break;
     case NOOP:
       taskMap.remove(s);
       break;
     case NO_JOB:
       GearmanTask preSleepTask =
           new GearmanTask(
               new GrabJobEventHandler(s),
               new GearmanPacketImpl(
                   GearmanPacketMagic.REQ, GearmanPacketType.PRE_SLEEP, new byte[0]));
       taskMap.put(s, preSleepTask);
       s.submitTask(preSleepTask);
       break;
     case ECHO_RES:
       break;
     case OPTION_RES:
       break;
     case ERROR:
       s.closeSession();
       break;
     default:
       LOG.warn(
           "Received unknown packet type " + t + " from session " + s + ". Closing connection.");
       s.closeSession();
   }
 }
  private GearmanJobStatus updateJobStatus(byte[] jobhandle, GearmanJobServerSession session)
      throws IOException, IllegalStateException, GearmanException {
    if (!runState.equals(state.RUNNING)) {
      throw new IllegalStateException(CLIENT_NOT_ACTIVE);
    }

    if (jobhandle == null || jobhandle.length == 0) {
      throw new IllegalStateException("Invalid job handle. Handle must" + " not be null nor empty");
    }
    GearmanPacket statusRequest =
        new GearmanPacketImpl(GearmanPacketMagic.REQ, GearmanPacketType.GET_STATUS, jobhandle);
    GearmanServerResponseHandler handler =
        (GearmanServerResponseHandler) new GearmanJobStatusImpl();
    GearmanTask t = new GearmanTask(handler, statusRequest);
    session.submitTask(t);
    if (!driveRequestTillState(t, GearmanTask.State.FINISHED)) {
      throw new GearmanException(
          "Failed to execute jobstatus request " + t + " to session " + session);
    }
    return (GearmanJobStatus) handler;
  }
  private boolean setForwardExceptions(GearmanJobServerSession session) {
    GearmanServerResponseHandler optionReqHandler =
        new GearmanServerResponseHandler() {
          boolean isDone = false;

          public boolean isDone() {
            return isDone;
          }

          public void handleEvent(GearmanPacket event) throws GearmanException {
            GearmanPacketType type = event.getPacketType();
            if (type.equals(GearmanPacketType.OPTION_RES)
                && Arrays.equals(EXCEPTIONS, event.getData())) {
              isDone = true;
            }
          }
        };
    GearmanTask setExceptionsTask =
        new GearmanTask(
            optionReqHandler,
            new GearmanPacketImpl(
                GearmanPacketMagic.REQ, GearmanPacketType.OPTION_REQ, EXCEPTIONS));
    session.submitTask(setExceptionsTask);
    Exception e = null;
    boolean success = false;
    try {
      success = driveRequestTillState(setExceptionsTask, GearmanTask.State.FINISHED);
    } catch (IOException ioe) {
      e = ioe;
    }
    if (!success) {
      if (e != null) {
        LOG.info("Failed to set forward-exceptions option to " + session.getConnection(), e);
      } else {
        LOG.info("Failed to set forward-exceptions option to " + session.getConnection());
      }
      return false;
    }
    return true;
  }
 public byte[] echo(byte[] data) throws IOException, GearmanException {
   if (!runState.equals(state.RUNNING)) {
     throw new IllegalStateException(CLIENT_NOT_ACTIVE);
   }
   GearmanPacket echoRequest =
       new GearmanPacketImpl(GearmanPacketMagic.REQ, GearmanPacketType.ECHO_REQ, data);
   GearmanEchoResponseHandler handler = new GearmanEchoResponseHandler();
   GearmanTask t = new GearmanTask(handler, echoRequest);
   GearmanJobServerSession session = getSessionForTask();
   session.submitTask(t);
   LOG.info(
       "Client "
           + this
           + " has submitted echo request "
           + "(payload = "
           + ByteUtils.toHex(data)
           + " to session "
           + session);
   if (!driveRequestTillState(t, GearmanTask.State.FINISHED)) {
     throw new GearmanException("Failed to execute echo request " + t + " to session " + session);
   }
   LOG.info("Client " + this + " has completed echo request " + "to session " + session);
   return handler.getResults();
 }
  public <T> Future<T> submit(Callable<T> task) {

    if (task == null) {
      throw new IllegalStateException("Null task was submitted to " + "gearman client");
    }

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

    if (!(task instanceof GearmanServerResponseHandler)) {
      throw new RejectedExecutionException(
          "Task must implement the "
              + GearmanServerResponseHandler.class
              + " interface to"
              + " submitted to this client");
    }

    GearmanJobImpl job = (GearmanJobImpl) task;
    GearmanServerResponseHandler handler = (GearmanServerResponseHandler) job;

    if (job.isDone()) {
      throw new RejectedExecutionException("Task can not be resubmitted ");
    }

    GearmanJobServerSession session = null;
    try {
      session = getSessionForTask();
    } catch (IOException ioe) {
      throw new RejectedExecutionException(ioe);
    }
    job.setJobServerSession(session);
    GearmanPacket submitRequest = getPacketFromJob(job);
    GearmanTask submittedJob = new GearmanTask(handler, submitRequest);
    session.submitTask(submittedJob);
    LOG.info(
        "Client "
            + this
            + " has submitted job "
            + job
            + // NOPMD
            " to session "
            + session
            + ". Job has been added to the "
            + // NOPMD
            "active job queue");
    try {
      jobAwatingCreation = job;
      if (!(driveRequestTillState(submittedJob, GearmanTask.State.RUNNING))) {
        throw new RejectedExecutionException(
            "Timed out waiting for" + " submission of " + job + " to complete");
      }
    } catch (IOException ioe) {
      LOG.warn("Client " + this + " encounted an " + "IOException while drivingIO", ioe);
      throw new RejectedExecutionException(
          "Failed to successfully " + "submit" + job + " due to IOException", ioe);
    } finally {
      jobAwatingCreation = null;
    }

    return (Future<T>) job;
  }
 private void sendToAll(GearmanServerResponseHandler handler, GearmanPacket p) {
   GearmanTask gsr = new GearmanTask(handler, p);
   for (GearmanJobServerSession sess : sessionMap.values()) {
     sess.submitTask(gsr);
   }
 }
  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);
  }