org.glassfish.jersey.server.ChunkedOutput Java Examples

/**
 * Handles the RESTCONF Event Notification Subscription request. If the
 * subscription is successful, a ChunkedOutput stream is created and returned
 * to the caller.
 * <p>
 * This function is not blocked on streaming the data (so that it can handle
 * other incoming requests). Instead, a worker thread running in the background
 * does the data streaming. If errors occur during streaming, the worker thread
 * calls ChunkedOutput.close() to disconnect the session and terminates itself.
 *
 * @param streamId Event stream ID
 * @param request  RESTCONF client information from which the client IP
 *                 address is retrieved
 * @return A string data stream over HTTP keep-alive session
 */
@GET
@Produces(MediaTypeRestconf.APPLICATION_YANG_DATA_JSON)
@Path("streams/{streamId}")
public ChunkedOutput<String> handleNotificationRegistration(@PathParam("streamId") String streamId,
                                                            @Context HttpServletRequest request) {
    final ChunkedOutput<String> output = new ChunkedOutput<>(String.class);
    try {
        service.subscribeEventStream(streamId, request.getRemoteAddr(), output);
    } catch (RestconfException e) {
        log.error("ERROR: handleNotificationRegistration: {}", e.getMessage());
        log.debug("Exception in handleNotificationRegistration:", e);
        try {
            output.close();
        } catch (IOException ex) {
            log.error("ERROR: handleNotificationRegistration:", ex);
        }
    }

    return output;
}
 

/**
 * Send operation to agent and get response as stream.
 *
 * @param operationRequest - Operation request.
 * @return The operation response.
 */
@POST
@Path("stream-operation")
@Consumes(MediaType.APPLICATION_JSON)
public ChunkedOutput<String> sendStreamingOperation(OperationRequest operationRequest) {
    final ChunkedOutput<String> streamingBuffer = new ChunkedOutput<String>(String.class);
    logger.info("Operation request received " + operationRequest.toJSON());
    AgentStreamHandler agentStreamHandler = new AgentStreamHandler(streamingBuffer, operationRequest,
            operationRequest.getAgentId());
    try {
        agentStreamHandler.startSendCommand();
        SessionManager.getAgentObservable().addObserver(agentStreamHandler);
    } catch (AgentHandleException e) {
        logger.error("Error while sending command to the Agent for agent " + operationRequest.getAgentId(), e);
        try {
            streamingBuffer.close();
        } catch (IOException ioe) {
            logger.error("Error while closing output stream for agent " + operationRequest.getAgentId(), ioe);
        }
    }
    return streamingBuffer;
}
 

/**
 * Agent handler constructor to initialize streaming object.
 *
 * @param streamingBuffer ChunkedOutput buffer to write back result
 * @param operationRequest Request from the test runner
 */
public AgentStreamHandler(ChunkedOutput<String> streamingBuffer, OperationRequest operationRequest,
                          String agentId) {
    this.streamingBuffer = streamingBuffer;
    this.operationRequest = operationRequest;
    this.agentId = agentId;
}
 

public ChunkedOutput<String> handleRequest(ExamplesIterable examplesIterable) {

		ChunkedOutput<String> chunkedOutput = new ChunkedOutput<String>(String.class);

		// get the example processor.
		ExampleProcessor exampleProcessor = exampleProcessorFactory.getExampleProcessor(examplesIterable);

		if (exampleProcessor.getExampleProcessorFeatures().isAsync() == false)
			submitSynchronously(exampleProcessor, chunkedOutput);
		else {
			submitAsynchronously(exampleProcessor, chunkedOutput);
		}

		return chunkedOutput;
	}
 

public ChunkedOutput<byte[]> startChunckedDownload(JobId previewJobId, String currentUser, DownloadFormat downloadFormat,
                                                   long delay) throws JobNotFoundException {

  //first check that current user has access to preview data
  final GetJobRequest previewJobRequest = GetJobRequest.newBuilder()
    .setJobId(previewJobId)
    .setUserName(currentUser)
    .build();

  // ensure that we could access to the job.
  final JobDetails previewJobDetails = jobsService.getJobDetails(JobDetailsRequest.newBuilder()
    .setJobId(JobsProtoUtil.toBuf(previewJobId))
    .setUserName(currentUser)
    .build());
  final JobInfo previewJobInfo = JobsProtoUtil.getLastAttempt(previewJobDetails).getInfo();
  final List<String> datasetPath = previewJobInfo.getDatasetPathList();

  if (!JOB_TYPES_TO_DOWNLOAD.contains(previewJobInfo.getQueryType())) {
    logger.error("Not supported job type: {} for job '{}'. Supported job types are: {}",
      previewJobInfo.getQueryType(), previewJobId,
      String.join(", ", JOB_TYPES_TO_DOWNLOAD.stream()
        .map(queryType -> String.valueOf(queryType.getNumber()))
        .collect(Collectors.toList())));
    throw new IllegalArgumentException("Data for the job could not be downloaded");
  }

  final ChunkedOutput<byte[]> output = new ChunkedOutput<>(byte[].class);

  executorService.schedule(() -> {
    getJobResults(previewJobId, datasetPath, downloadFormat, currentUser, output);
  }, delay, TimeUnit.MILLISECONDS);

  return output;
}
 

@POST
@Consumes({ ExampleMediaTypes.PLAINTEXT_0_1_0, MediaType.TEXT_PLAIN, ExampleMediaTypes.SIMPLE_PROTOBUF_0_1_0, ExampleMediaTypes.SIMPLE_JSON_0_1_0, ExampleMediaTypes.STRUCTURED_JSON_0_1_0 })
@Produces({ PredictionMediaTypes.PLAINTEXT_0_1_0 })
@Path("/main")
public ChunkedOutput<String> doPredict(ExamplesIterable examplesIterable) throws IOException {

	return new RequestHandler(executorService, exampleProcessorFactory).handleRequest(examplesIterable);
}
 

@Test
public void streamChunkedResponseTest() throws IOException {
	String channelCode = "mock_collection";
	String callbackName = null;
	Float rate = new Float(-1);
	String duration = "-1";
	ChunkedOutput<String> output = asyncStream.streamChunkedResponse(channelCode, callbackName, rate, duration);
	//
	Class<?> s = output.getRawType();
	assertNotNull(output);
	assertEquals("class java.lang.String", s.toString());
}
 

public AsyncStreamRedisSubscriber(final Jedis jedis, final ChunkedOutput<String> responseWriter,
		ArrayList<ChunkedOutput<String>> writerList,
		final SubscriptionDataObject subData) throws IOException {
	this.channel = subData.redisChannel;
	this.callbackName = subData.callbackName;
	this.responseWriter = responseWriter;
	this.writerList = writerList;

	this.subData = new SubscriptionDataObject();
	this.subData.set(subData);

	this.setRunFlag(true);		
	if (subData.duration != null) {
		subscriptionDuration = parseTime(subData.duration);
	} else {
		subscriptionDuration = SUBSCRIPTION_MAX_DURATION;
	}
	//System.out.println("rate=" + subData.rate + ", duration=" + subData.duration + ", callbackName=" + subData.callbackName);
	logger.info("Client requested subscription for duration = " + subscriptionDuration);
	if (subData.rate > 0) {
		messageRate = subData.rate;			// specified as messages/min (NOTE: upper-bound)
		sleepTime = Math.max(0, Math.round(60 * 1000 / messageRate));		// time to sleep between sends (in msecs)
	} else {
		sleepTime = DEFAULT_SLEEP_TIME;		// use default value
	}
	messageList = new ArrayList<String>(DEFAULT_COUNT);
}
 

@Override
public void contextInitialized(ServletContextEvent sce) {
	// Now initialize shared jedis connection object and thread pool object
	jedisConn = new JedisConnectionObject(redisHost, redisPort);
	executorService = Executors.newCachedThreadPool();
	writerList = new ArrayList<ChunkedOutput<String>>();
	logger.info("Context Initialized, executorService: " + executorService);
}
 

@Override
public void onClose(ChunkedOutput<OutboundEvent> event) {
    if (event instanceof SitemapEventOutput) {
        SitemapEventOutput sitemapEvent = (SitemapEventOutput) event;
        logger.debug("SSE connection for subscription {} has been closed.", sitemapEvent.getSubscriptionId());
        subscriptions.removeSubscription(sitemapEvent.getSubscriptionId());
        EventOutput eventOutput = eventOutputs.remove(sitemapEvent.getSubscriptionId());
        if (eventOutput != null) {
            broadcaster.remove(eventOutput);
        }
    }
}
 

@Override
public void subscribeEventStream(String streamId, String clientIpAddr,
                                 ChunkedOutput<String> output)
        throws RestconfException {
    //TODO: to be completed
    throw new RestconfException("Not implemented",
            RestconfError.ErrorTag.OPERATION_NOT_SUPPORTED,
            Response.Status.NOT_IMPLEMENTED,
            Optional.empty(), Optional.of("subscribeEventStream not yet implemented"));
}
 

/**
 * {@inheritDoc}
 */
@Override
public ChunkedOutput apply(ContainerRequestContext containerRequestContext) {
    return null;
}
 

/**
 * {@inheritDoc}
 */
@Override
public ChunkedOutput apply(ContainerRequestContext containerRequestContext) {
    return null;
}
 

/**
 * {@inheritDoc}
 */
@Override
public ChunkedOutput apply(ContainerRequestContext containerRequestContext) {
    return null;
}
 

/**
 * {@inheritDoc}
 */
@Override
public ChunkedOutput apply(ContainerRequestContext containerRequestContext) {
    return null;
}
 

/**
 * {@inheritDoc}
 */
@Override
public ChunkedOutput apply(ContainerRequestContext containerRequestContext) {
    return null;
}
 

/**
 * {@inheritDoc}
 */
@Override
public ChunkedOutput apply(ContainerRequestContext containerRequestContext) {
    return null;
}
 

@Override
public void onException(ChunkedOutput<OutboundEvent> event, Exception e) {
    // the exception is usually "null" and onClose() is automatically called afterwards
    // - so let's don't do anything in this method.
}
 

@Override
public void onClose(ChunkedOutput<OutboundEvent> chunkedOutput) {
  super.onClose(chunkedOutput);
  LOG.debug("[{}] SSE channel closed", this);
  outputs.decrementAndGet();
}
 

@Override
public <OUT extends ChunkedOutput<OutboundEvent>> boolean add(OUT chunkedOutput) {
  LOG.debug("[{}] Adding SSE channel", this, chunkedOutput);
  outputs.incrementAndGet();
  return super.add(chunkedOutput);
}
 

@Override
public void onException(ChunkedOutput<OutboundEvent> chunkedOutput, Exception exception) {
  super.onException(chunkedOutput, exception);
  LOG.debug("[{}] SSE exception", this);
  //closeJMX();
}
 

private void getJobResults(JobId previewJobId, List<String> datasetPath, DownloadFormat downloadFormat, String currentUser, ChunkedOutput<byte[]> output) {
  try {
    JobDataClientUtils.waitForFinalState(jobsService, previewJobId);
    JobDetails previewJobDetails = jobsService.getJobDetails(JobDetailsRequest.newBuilder()
      .setJobId(JobsProtoUtil.toBuf(previewJobId))
      .setUserName(currentUser)
      .build());
    //read current job state
    checkJobCompletionState(previewJobDetails);

    JobDetails downloadJobDetails = previewJobDetails;
    final JobInfo previewJobInfo = JobsProtoUtil.getLastAttempt(previewJobDetails).getInfo();

    if (previewJobInfo.getQueryType() != QueryType.UI_EXPORT) {
      DatasetDownloadManager manager = datasetService.downloadManager();

      JobId downloadJobId = manager.scheduleDownload(datasetPath, previewJobInfo.getSql(),
        downloadFormat, previewJobInfo.getContextList(), -1 /* no limit */, currentUser, previewJobId);

      JobDataClientUtils.waitForFinalState(jobsService, downloadJobId);
      //get the final state of a job after completion
      JobDetailsRequest jobDetailsRequest = JobDetailsRequest.newBuilder()
        .setUserName(currentUser)
        .setJobId(JobsProtoUtil.toBuf(downloadJobId))
        .build();
      downloadJobDetails = jobsService.getJobDetails(jobDetailsRequest);
      checkJobCompletionState(downloadJobDetails);
    }

    final DatasetDownloadManager.DownloadDataResponse downloadDataResponse =
      datasetService.downloadData(JobsProtoUtil.getLastAttempt(downloadJobDetails).getInfo().getDownloadInfo(), currentUser);

    try (InputStream input = downloadDataResponse.getInput(); ChunkedOutput toClose = output) {
      byte[] buf = new byte[4096];
      int bytesRead = input.read(buf);
      while (bytesRead >= 0) {
        if (bytesRead < buf.length) {
          output.write(Arrays.copyOf(buf, bytesRead));
        } else {
          output.write(buf);
        }
        bytesRead = input.read(buf);
      }
    }
  } catch (JobNotFoundException | IOException e) {
    throw new WebApplicationException(e);
  }
}
 

protected Response doDownload(JobId previewJobId, DownloadFormat downloadFormat) throws JobResourceNotFoundException, JobNotFoundException {
  final String currentUser = securityContext.getUserPrincipal().getName();

  //first check that current user has access to preview data
  final JobDetailsRequest previewJobRequest = JobDetailsRequest.newBuilder()
    .setJobId(JobsProtoUtil.toBuf(previewJobId))
    .setUserName(currentUser)
    .build();

  final JobDetails jobDetails = jobsService.getJobDetails(previewJobRequest);

  final DownloadUtil downloadUtil = new DownloadUtil(jobsService, datasetService);
  final ChunkedOutput<byte[]> output = downloadUtil.startChunckedDownload(previewJobId, currentUser, downloadFormat, getDelay());

  final String contentType;
  if (downloadFormat != null) {
    switch (downloadFormat) {
      case JSON:
        contentType = APPLICATION_JSON;
        break;
      case CSV:
        contentType = "text/csv";
        break;
      default:
        contentType = MediaType.APPLICATION_OCTET_STREAM;
        break;
    }
  } else {
    contentType = MediaType.APPLICATION_OCTET_STREAM;
  }

  final JobInfo info = JobsProtoUtil.getLastAttempt(jobDetails).getInfo();

  final String outputFileName;
  // job id is already a download job. So just extract a filename from it
  if (info.getQueryType() == QueryType.UI_EXPORT) {
    outputFileName = info.getDownloadInfo().getFileName();
  } else {
    // must use DatasetDownloadManager.getDownloadFileName. If a naming convention is changed, the change should go to
    // DatasetDownloadManager.getDownloadFileName method.
    outputFileName = DatasetDownloadManager.getDownloadFileName(previewJobId, downloadFormat);
  }

  return Response.ok(output, contentType)
    .header("Content-Disposition", "attachment; filename=\"" + outputFileName + "\"")
    // stops the browser from trying to determine the type of the file based on the content.
    .header( "X-Content-Type-Options", "nosniff")
    .build();
}
 

private void submitAsynchronously(final ExampleProcessor exampleSubmitter, final ChunkedOutput<String> chunkedOutput) {

		executorService.submit(new Runnable() {

			@Override
			public void run() {

				submitSynchronously(exampleSubmitter, chunkedOutput);

			}

		});

	}
 

/**
 * Handles an Event Stream subscription request. This function creates
 * a worker thread to listen to events and writes to a ChunkedOutput,
 * which is passed in from the caller. (The worker thread blocks if
 * no events arrive.) The ChuckedOutput is a pipe to which this
 * function acts as the writer and the caller the reader.
 * <p>
 * If the Event Stream cannot be subscribed due to reasons such as
 * the nonexistence of the target stream or failure to allocate
 * worker thread to handle the request, a RestconfException exception
 * is raised. The proper HTTP error status code is enclosed in the
 * exception, so that the caller may return it to the RESTCONF client
 * to display.
 *
 * @param streamId     ID of the RESTCONF stream to subscribe
 * @param clientIpAddr IP address of the RESTCONF client sending the request
 * @param output       A string data stream
 * @throws RestconfException if the Event Stream cannot be subscribed
 */
void subscribeEventStream(String streamId,
                          String clientIpAddr,
                          ChunkedOutput<String> output)
        throws RestconfException;