@Override protected Runnable wrapTask(Runnable command) { TraceToken token = getCurrentTraceToken(); return () -> { try (TraceTokenScope ignored = registerTraceToken(token)) { command.run(); } }; } }
/** * Create and register a new trace token with a random id. * @param properties Additional properties to include in the token. * @return The id of the created token. */ public static String createAndRegisterNewRequestToken(String... properties) { checkArgument((properties.length % 2) == 0, "odd number of elements in properties"); String newToken = UUID.randomUUID().toString(); registerRequestToken(newToken); if (properties.length != 0) { addTraceTokenProperties(properties); } return newToken; }
@Override public void close() { registerTraceToken(oldToken); } }
@Override public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain chain) throws IOException, ServletException { HttpServletRequest request = (HttpServletRequest) servletRequest; HttpServletResponse response = (HttpServletResponse) servletResponse; String token = request.getHeader("X-Proofpoint-TraceToken"); if (token == null || token.isEmpty()) { registerNewRequestToken(request); } else if (token.charAt(0) == '{') { try { registerTraceToken(TRACE_TOKEN_JSON_CODEC.fromJson(token)); } catch (RuntimeException e) { registerNewRequestToken(request); } } else { registerRequestToken(token); } chain.doFilter(request, response); }
DynamicBodySourceContentProvider(DynamicBodySource dynamicBodySource, AtomicLong bytesWritten) { this.dynamicBodySource = dynamicBodySource; this.bytesWritten = bytesWritten; traceToken = getCurrentTraceToken(); }
@Override public Request filterRequest(Request request) { checkNotNull(request, "request is null"); String token = traceTokenManager.getCurrentRequestToken(); if (token == null) { return request; } return fromRequest(request) .addHeader(TRACETOKEN_HEADER, token) .build(); }
private void registerNewRequestToken(HttpServletRequest request) throws UnknownHostException { byte[] randomBytes = new byte[15]; SECURE_RANDOM.get().nextBytes(randomBytes); registerRequestToken(tokenPrefix + encodeAddress(InetAddress.getByName(clientAddressExtractor.clientAddressFor(request))) + BASE64_URL_ENCODER.encodeToString(randomBytes) ); } }
@Override public void close() { if (writer instanceof AutoCloseable) { try (TraceTokenScope ignored = registerTraceToken(traceToken)) { ((AutoCloseable)writer).close(); } catch (Exception e) { throwIfUnchecked(e); throw new RuntimeException(e); } } } }
JettyResponseFuture(JettyHttpClient jettyHttpClient, Request request, org.eclipse.jetty.client.api.Request jettyRequest, ResponseHandler<T, E> responseHandler, AtomicLong bytesWritten, RequestStats stats) { this.jettyHttpClient = requireNonNull(jettyHttpClient, "jettyHttpClientf is null"); this.request = requireNonNull(request, "request is null"); this.jettyRequest = requireNonNull(jettyRequest, "jettyRequest is null"); this.responseHandler = requireNonNull(responseHandler, "responseHandler is null"); this.bytesWritten = requireNonNull(bytesWritten, "bytesWritten is null"); this.stats = requireNonNull(stats, "stats is null"); traceToken = getCurrentTraceToken(); }
String token = getCurrentRequestToken();
@Override protected <T> Callable<T> wrapTask(Callable<T> callable) { TraceToken token = getCurrentTraceToken(); return () -> { try (TraceTokenScope ignored = registerTraceToken(token)) { return callable.call(); } }; }
/** * Associate a given trace token id, with no other properties, with the * current thread. * * @param tokenId The tokenId to associate with the current thread, or null to * remove the thread's token. * @return a {@link TraceTokenScope} which may be used to restore the thread's * previous association. Intended to be used with try-with-resources: * <code> * try (TraceTokenScope ignored = registerRequestToken(traceTokenId)) { * // process request * } * </code> */ public static TraceTokenScope registerRequestToken(@Nullable String tokenId) { if (tokenId == null) { return registerTraceToken(null); } return registerTraceToken(new TraceToken(ImmutableMap.of("id", tokenId))); }
@Override public <T> ListenableFuture<Void> post(final Iterable<T> events) throws IllegalArgumentException { requireNonNull(events, "eventsSupplier is null"); TraceToken token = getCurrentTraceToken(); Request request = preparePost() .setUri(URI.create("v2/event")) .setHeader("User-Agent", nodeInfo.getNodeId()) .setHeader("Content-Type", MEDIA_TYPE_JSON.toString()) .setBodySource((DynamicBodySource) out -> eventWriter.createEventWriter(events.iterator(), token, out)) .build(); return httpClient.executeAsync(request, EVENT_RESPONSE_HANDLER); }
@Override public <T> CheckedFuture<Void, RuntimeException> post(EventGenerator<T> eventGenerator) { checkNotNull(eventGenerator, "eventGenerator is null"); String token = traceTokenManager.getCurrentRequestToken(); List<URI> uris = serviceSelector.selectHttpService(); if (uris.isEmpty()) { return Futures.<Void, RuntimeException>immediateFailedCheckedFuture(new ServiceUnavailableException(serviceSelector.getType(), serviceSelector.getPool())); } // todo this doesn't really work due to returning the future which can fail without being retried Request request = preparePost() .setUri(uris.get(0).resolve("/v2/event")) .setHeader("User-Agent", nodeInfo.getNodeId()) .setHeader("Content-Type", MEDIA_TYPE_JSON.toString()) .setBodyGenerator(new JsonEntityWriter<>(eventWriter, eventGenerator, token)) .build(); return httpClient.executeAsync(request, new EventResponseHandler(serviceSelector.getType(), serviceSelector.getPool())); }
TraceToken traceToken = getCurrentTraceToken(); synchronized (subFutureLock) { Duration backoff = attemptBackoffPolicy.backoff(previousBackoff); ScheduledFuture<?> scheduledFuture = retryExecutor.schedule(() -> { try (TraceTokenScope scope = registerTraceToken(traceToken)){ synchronized (subFutureLock) { HttpServiceAttempt nextAttempt;
/** * Add properties to the current thread's trace token. If there is * currently no trace token, does nothing. * * @param properties Properties to add or replace. * @return a {@link TraceTokenScope} which may be used to restore the thread's * previous set of properties. */ public static TraceTokenScope addTraceTokenProperties(String... properties) { TokenState tokenState = token.get(); if (tokenState == null) { return new TraceTokenScope(null); } Map<String, String> map = new LinkedHashMap<>(tokenState.getToken()); checkArgument((properties.length % 2) == 0, "odd number of elements in properties"); for (int i = 0; i < properties.length; i += 2) { requireNonNull(properties[i], "property key is null"); requireNonNull(properties[i+1], "property value is null"); map.put(properties[i], properties[i+1]); } return registerTraceToken(new TraceToken(map)); }
@Override public Request filterRequest(Request request) { requireNonNull(request, "request is null"); TraceToken token = getCurrentTraceToken(); if (token == null) { return request; } String tokenString; if (token.size() == 1) { tokenString = token.toString(); } else { tokenString = TRACE_TOKEN_JSON_CODEC.toJson(token); } return fromRequest(request) .addHeader(TRACETOKEN_HEADER, tokenString) .build(); }
@Override public boolean cancel(boolean mayInterruptIfRunning) { try { state.set(JettyAsyncHttpState.CANCELED); jettyRequest.abort(new CancellationException()); return super.cancel(mayInterruptIfRunning); } catch (Throwable e) { try (TraceTokenScope ignored = registerTraceToken(traceToken)) { setException(e); } return true; } }
@Override @SuppressWarnings("ReferenceEquality") // Reference equality to DONE is intentional protected ByteBuffer computeNext() { ByteBuffer chunk = chunks.poll(); if (chunk == null) { try (TraceTokenScope ignored = registerTraceToken(traceToken)) { while (chunk == null) { try { writer.write(); } catch (Exception e) { throwIfUnchecked(e); throw new RuntimeException(e); } chunk = chunks.poll(); } } } if (chunk == DONE) { return endOfData(); } bytesWritten.addAndGet(chunk.position()); ((Buffer)chunk).flip(); return chunk; }
void completed(Response response, InputStream content) { if (state.get() == JettyAsyncHttpState.CANCELED) { return; } try (TraceTokenScope ignored = registerTraceToken(traceToken)) { T value; try { value = processResponse(response, content); } catch (Throwable e) { // this will be an instance of E from the response handler or an Error storeException(e); return; } state.set(JettyAsyncHttpState.DONE); set(value); } }