def _set_state(self, shard_state, tstate, task_directive):
    """Set shard_state and tstate based on task_directive.

    Args:
      shard_state: model.ShardState for current shard.
      tstate: model.TransientShardState for current shard.
      task_directive: self._TASK_DIRECTIVE for current shard.

    Returns:
      A _TASK_DIRECTIVE enum.
      PROCEED_TASK if task should proceed normally.
      RETRY_SHARD if shard should be retried.
      RETRY_SLICE if slice should be retried.
      FAIL_TASK if sahrd should fail.
      RECOVER_SLICE if slice should be recovered.
      ABORT_SHARD if shard should be aborted.
      RETRY_TASK if task should be retried.
      DROP_TASK if task should be dropped.
    """
    if task_directive in (self._TASK_DIRECTIVE.RETRY_TASK,
                          self._TASK_DIRECTIVE.DROP_TASK):
      return task_directive

    if task_directive == self._TASK_DIRECTIVE.ABORT_SHARD:
      shard_state.set_for_abort()
      return task_directive

    if task_directive == self._TASK_DIRECTIVE.PROCEED_TASK:
      shard_state.advance_for_next_slice()
      tstate.advance_for_next_slice()
      return task_directive

    if task_directive == self._TASK_DIRECTIVE.RECOVER_SLICE:
      tstate.advance_for_next_slice(recovery_slice=True)
      shard_state.advance_for_next_slice(recovery_slice=True)
      return task_directive

    if task_directive == self._TASK_DIRECTIVE.RETRY_SLICE:
      task_directive = self._attempt_slice_retry(shard_state, tstate)
    if task_directive == self._TASK_DIRECTIVE.RETRY_SHARD:
      task_directive = self._attempt_shard_retry(shard_state, tstate)
    if task_directive == self._TASK_DIRECTIVE.FAIL_TASK:
      shard_state.set_for_failure()

    return task_directive