예제 #1
0
  @Override
  public Sequence<T> run(final Query<T> query) {
    final QueryToolChest<T, Query<T>> toolChest = warehouse.getToolChest(query);
    final CacheStrategy<T, Object, Query<T>> strategy = toolChest.getCacheStrategy(query);

    final Map<DruidServer, List<SegmentDescriptor>> serverSegments = Maps.newTreeMap();

    final List<Pair<DateTime, byte[]>> cachedResults = Lists.newArrayList();
    final Map<String, CachePopulator> cachePopulatorMap = Maps.newHashMap();

    final boolean useCache =
        Boolean.parseBoolean(query.getContextValue("useCache", "true")) && strategy != null;
    final boolean populateCache =
        Boolean.parseBoolean(query.getContextValue("populateCache", "true")) && strategy != null;
    final boolean isBySegment = Boolean.parseBoolean(query.getContextValue("bySegment", "false"));

    ImmutableMap.Builder<String, String> contextBuilder = new ImmutableMap.Builder<>();

    final String priority = query.getContextValue("priority", "0");
    contextBuilder.put("priority", priority);

    if (populateCache) {
      contextBuilder.put("bySegment", "true");
    }
    contextBuilder.put("intermediate", "true");

    final Query<T> rewrittenQuery = query.withOverriddenContext(contextBuilder.build());

    VersionedIntervalTimeline<String, ServerSelector> timeline =
        serverView.getTimeline(query.getDataSource());
    if (timeline == null) {
      return Sequences.empty();
    }

    // build set of segments to query
    Set<Pair<ServerSelector, SegmentDescriptor>> segments = Sets.newLinkedHashSet();

    List<TimelineObjectHolder<String, ServerSelector>> serversLookup = Lists.newLinkedList();

    for (Interval interval : rewrittenQuery.getIntervals()) {
      serversLookup.addAll(timeline.lookup(interval));
    }

    // Let tool chest filter out unneeded segments
    final List<TimelineObjectHolder<String, ServerSelector>> filteredServersLookup =
        toolChest.filterSegments(query, serversLookup);

    for (TimelineObjectHolder<String, ServerSelector> holder : filteredServersLookup) {
      for (PartitionChunk<ServerSelector> chunk : holder.getObject()) {
        ServerSelector selector = chunk.getObject();
        final SegmentDescriptor descriptor =
            new SegmentDescriptor(
                holder.getInterval(), holder.getVersion(), chunk.getChunkNumber());

        segments.add(Pair.of(selector, descriptor));
      }
    }

    final byte[] queryCacheKey;
    if (strategy != null) {
      queryCacheKey = strategy.computeCacheKey(query);
    } else {
      queryCacheKey = null;
    }

    if (queryCacheKey != null) {
      Map<Pair<ServerSelector, SegmentDescriptor>, Cache.NamedKey> cacheKeys = Maps.newHashMap();
      for (Pair<ServerSelector, SegmentDescriptor> segment : segments) {
        final Cache.NamedKey segmentCacheKey =
            computeSegmentCacheKey(
                segment.lhs.getSegment().getIdentifier(), segment.rhs, queryCacheKey);
        cacheKeys.put(segment, segmentCacheKey);
      }

      // Pull cached segments from cache and remove from set of segments to query
      final Map<Cache.NamedKey, byte[]> cachedValues;
      if (useCache) {
        cachedValues = cache.getBulk(cacheKeys.values());
      } else {
        cachedValues = ImmutableMap.of();
      }

      for (Map.Entry<Pair<ServerSelector, SegmentDescriptor>, Cache.NamedKey> entry :
          cacheKeys.entrySet()) {
        Pair<ServerSelector, SegmentDescriptor> segment = entry.getKey();
        Cache.NamedKey segmentCacheKey = entry.getValue();
        final Interval segmentQueryInterval = segment.rhs.getInterval();

        final byte[] cachedValue = cachedValues.get(segmentCacheKey);
        if (cachedValue != null) {
          // remove cached segment from set of segments to query
          segments.remove(segment);
          cachedResults.add(Pair.of(segmentQueryInterval.getStart(), cachedValue));
        } else if (populateCache) {
          final String segmentIdentifier = segment.lhs.getSegment().getIdentifier();
          cachePopulatorMap.put(
              String.format("%s_%s", segmentIdentifier, segmentQueryInterval),
              new CachePopulator(cache, objectMapper, segmentCacheKey));
        }
      }
    }

    // Compile list of all segments not pulled from cache
    for (Pair<ServerSelector, SegmentDescriptor> segment : segments) {
      final QueryableDruidServer queryableDruidServer = segment.lhs.pick();

      if (queryableDruidServer == null) {
        log.error("No servers found for %s?! How can this be?!", segment.rhs);
      } else {
        final DruidServer server = queryableDruidServer.getServer();
        List<SegmentDescriptor> descriptors = serverSegments.get(server);

        if (descriptors == null) {
          descriptors = Lists.newArrayList();
          serverSegments.put(server, descriptors);
        }

        descriptors.add(segment.rhs);
      }
    }

    return new LazySequence<>(
        new Supplier<Sequence<T>>() {
          @Override
          public Sequence<T> get() {
            ArrayList<Pair<DateTime, Sequence<T>>> listOfSequences = Lists.newArrayList();

            addSequencesFromServer(listOfSequences);
            addSequencesFromCache(listOfSequences);

            Collections.sort(
                listOfSequences,
                Ordering.natural().onResultOf(Pair.<DateTime, Sequence<T>>lhsFn()));

            final Sequence<Sequence<T>> seq =
                Sequences.simple(
                    Iterables.transform(listOfSequences, Pair.<DateTime, Sequence<T>>rhsFn()));
            if (strategy == null) {
              return toolChest.mergeSequences(seq);
            } else {
              return strategy.mergeSequences(seq);
            }
          }

          private void addSequencesFromCache(
              ArrayList<Pair<DateTime, Sequence<T>>> listOfSequences) {
            if (strategy == null) {
              return;
            }

            final Function<Object, T> pullFromCacheFunction = strategy.pullFromCache();
            final TypeReference<Object> cacheObjectClazz = strategy.getCacheObjectClazz();
            for (Pair<DateTime, byte[]> cachedResultPair : cachedResults) {
              final byte[] cachedResult = cachedResultPair.rhs;
              Sequence<Object> cachedSequence =
                  new BaseSequence<>(
                      new BaseSequence.IteratorMaker<Object, Iterator<Object>>() {
                        @Override
                        public Iterator<Object> make() {
                          try {
                            if (cachedResult.length == 0) {
                              return Iterators.emptyIterator();
                            }

                            return objectMapper.readValues(
                                objectMapper.getFactory().createParser(cachedResult),
                                cacheObjectClazz);
                          } catch (IOException e) {
                            throw Throwables.propagate(e);
                          }
                        }

                        @Override
                        public void cleanup(Iterator<Object> iterFromMake) {}
                      });
              listOfSequences.add(
                  Pair.of(
                      cachedResultPair.lhs, Sequences.map(cachedSequence, pullFromCacheFunction)));
            }
          }

          @SuppressWarnings("unchecked")
          private void addSequencesFromServer(
              ArrayList<Pair<DateTime, Sequence<T>>> listOfSequences) {
            for (Map.Entry<DruidServer, List<SegmentDescriptor>> entry :
                serverSegments.entrySet()) {
              final DruidServer server = entry.getKey();
              final List<SegmentDescriptor> descriptors = entry.getValue();

              final QueryRunner clientQueryable = serverView.getQueryRunner(server);
              if (clientQueryable == null) {
                log.makeAlert("WTF!? server[%s] doesn't have a client Queryable?", server).emit();
                continue;
              }

              final Sequence<T> resultSeqToAdd;
              final MultipleSpecificSegmentSpec segmentSpec =
                  new MultipleSpecificSegmentSpec(descriptors);
              List<Interval> intervals = segmentSpec.getIntervals();

              if (!server.isAssignable() || !populateCache || isBySegment) {
                resultSeqToAdd = clientQueryable.run(query.withQuerySegmentSpec(segmentSpec));
              } else {
                resultSeqToAdd =
                    toolChest.mergeSequences(
                        Sequences.map(
                            clientQueryable.run(rewrittenQuery.withQuerySegmentSpec(segmentSpec)),
                            new Function<Object, Sequence<T>>() {
                              private final Function<T, Object> prepareForCache =
                                  strategy.prepareForCache();

                              @Override
                              public Sequence<T> apply(Object input) {
                                Result<Object> result = (Result<Object>) input;
                                final BySegmentResultValueClass<T> value =
                                    (BySegmentResultValueClass<T>) result.getValue();
                                String segmentIdentifier = value.getSegmentId();
                                final Iterable<T> segmentResults = value.getResults();

                                CachePopulator cachePopulator =
                                    cachePopulatorMap.get(
                                        String.format(
                                            "%s_%s", segmentIdentifier, value.getInterval()));
                                if (cachePopulator != null) {
                                  cachePopulator.populate(
                                      Iterables.transform(segmentResults, prepareForCache));
                                }

                                return Sequences.simple(
                                    Iterables.transform(
                                        segmentResults,
                                        toolChest.makeMetricManipulatorFn(
                                            rewrittenQuery,
                                            new MetricManipulationFn() {
                                              @Override
                                              public Object manipulate(
                                                  AggregatorFactory factory, Object object) {
                                                return factory.deserialize(object);
                                              }
                                            })));
                              }
                            }));
              }

              listOfSequences.add(Pair.of(intervals.get(0).getStart(), resultSeqToAdd));
            }
          }
        });
  }
예제 #2
0
  @POST
  @Produces({MediaType.APPLICATION_JSON, SmileMediaTypes.APPLICATION_JACKSON_SMILE})
  @Consumes({
    MediaType.APPLICATION_JSON,
    SmileMediaTypes.APPLICATION_JACKSON_SMILE,
    APPLICATION_SMILE
  })
  public Response doPost(
      InputStream in,
      @QueryParam("pretty") String pretty,
      @Context
          final HttpServletRequest req // used only to get request content-type and remote address
      ) throws IOException {
    final long start = System.currentTimeMillis();
    Query query = null;
    String queryId = null;

    final String reqContentType = req.getContentType();
    final boolean isSmile =
        SmileMediaTypes.APPLICATION_JACKSON_SMILE.equals(reqContentType)
            || APPLICATION_SMILE.equals(reqContentType);
    final String contentType =
        isSmile ? SmileMediaTypes.APPLICATION_JACKSON_SMILE : MediaType.APPLICATION_JSON;

    ObjectMapper objectMapper = isSmile ? smileMapper : jsonMapper;
    final ObjectWriter jsonWriter =
        pretty != null ? objectMapper.writerWithDefaultPrettyPrinter() : objectMapper.writer();

    final String currThreadName = Thread.currentThread().getName();
    try {
      query = objectMapper.readValue(in, Query.class);
      queryId = query.getId();
      if (queryId == null) {
        queryId = UUID.randomUUID().toString();
        query = query.withId(queryId);
      }
      if (query.getContextValue(QueryContextKeys.TIMEOUT) == null) {
        query =
            query.withOverriddenContext(
                ImmutableMap.of(
                    QueryContextKeys.TIMEOUT,
                    config.getMaxIdleTime().toStandardDuration().getMillis()));
      }

      Thread.currentThread()
          .setName(
              String.format(
                  "%s[%s_%s_%s]", currThreadName, query.getType(), query.getDataSource(), queryId));
      if (log.isDebugEnabled()) {
        log.debug("Got query [%s]", query);
      }

      final Map<String, Object> responseContext = new MapMaker().makeMap();
      final Sequence res = query.run(texasRanger, responseContext);
      final Sequence results;
      if (res == null) {
        results = Sequences.empty();
      } else {
        results = res;
      }

      final Yielder yielder =
          results.toYielder(
              null,
              new YieldingAccumulator() {
                @Override
                public Object accumulate(Object accumulated, Object in) {
                  yield();
                  return in;
                }
              });

      try {
        final Query theQuery = query;
        Response.ResponseBuilder builder =
            Response.ok(
                    new StreamingOutput() {
                      @Override
                      public void write(OutputStream outputStream)
                          throws IOException, WebApplicationException {
                        // json serializer will always close the yielder
                        CountingOutputStream os = new CountingOutputStream(outputStream);
                        jsonWriter.writeValue(os, yielder);

                        os
                            .flush(); // Some types of OutputStream suppress flush errors in the
                                      // .close() method.
                        os.close();

                        final long queryTime = System.currentTimeMillis() - start;
                        emitter.emit(
                            DruidMetrics.makeQueryTimeMetric(
                                    jsonMapper, theQuery, req.getRemoteAddr())
                                .setDimension("success", "true")
                                .build("query/time", queryTime));
                        emitter.emit(
                            DruidMetrics.makeQueryTimeMetric(
                                    jsonMapper, theQuery, req.getRemoteAddr())
                                .build("query/bytes", os.getCount()));

                        requestLogger.log(
                            new RequestLogLine(
                                new DateTime(start),
                                req.getRemoteAddr(),
                                theQuery,
                                new QueryStats(
                                    ImmutableMap.<String, Object>of(
                                        "query/time",
                                        queryTime,
                                        "query/bytes",
                                        os.getCount(),
                                        "success",
                                        true))));
                      }
                    },
                    contentType)
                .header("X-Druid-Query-Id", queryId);

        // Limit the response-context header, see https://github.com/druid-io/druid/issues/2331
        // Note that Response.ResponseBuilder.header(String key,Object value).build() calls
        // value.toString()
        // and encodes the string using ASCII, so 1 char is = 1 byte
        String responseCtxString = jsonMapper.writeValueAsString(responseContext);
        if (responseCtxString.length() > RESPONSE_CTX_HEADER_LEN_LIMIT) {
          log.warn(
              "Response Context truncated for id [%s] . Full context is [%s].",
              queryId, responseCtxString);
          responseCtxString = responseCtxString.substring(0, RESPONSE_CTX_HEADER_LEN_LIMIT);
        }

        return builder.header("X-Druid-Response-Context", responseCtxString).build();
      } catch (Exception e) {
        // make sure to close yielder if anything happened before starting to serialize the
        // response.
        yielder.close();
        throw Throwables.propagate(e);
      } finally {
        // do not close yielder here, since we do not want to close the yielder prior to
        // StreamingOutput having iterated over all the results
      }
    } catch (QueryInterruptedException e) {
      try {
        log.info("%s [%s]", e.getMessage(), queryId);
        final long queryTime = System.currentTimeMillis() - start;
        emitter.emit(
            DruidMetrics.makeQueryTimeMetric(jsonMapper, query, req.getRemoteAddr())
                .setDimension("success", "false")
                .build("query/time", queryTime));
        requestLogger.log(
            new RequestLogLine(
                new DateTime(start),
                req.getRemoteAddr(),
                query,
                new QueryStats(
                    ImmutableMap.<String, Object>of(
                        "query/time",
                        queryTime,
                        "success",
                        false,
                        "interrupted",
                        true,
                        "reason",
                        e.toString()))));
      } catch (Exception e2) {
        log.error(e2, "Unable to log query [%s]!", query);
      }
      return Response.serverError()
          .type(contentType)
          .entity(
              jsonWriter.writeValueAsBytes(
                  ImmutableMap.of(
                      "error", e.getMessage() == null ? "null exception" : e.getMessage())))
          .build();
    } catch (Exception e) {
      // Input stream has already been consumed by the json object mapper if query == null
      final String queryString = query == null ? "unparsable query" : query.toString();

      log.warn(e, "Exception occurred on request [%s]", queryString);

      try {
        final long queryTime = System.currentTimeMillis() - start;
        emitter.emit(
            DruidMetrics.makeQueryTimeMetric(jsonMapper, query, req.getRemoteAddr())
                .setDimension("success", "false")
                .build("query/time", queryTime));
        requestLogger.log(
            new RequestLogLine(
                new DateTime(start),
                req.getRemoteAddr(),
                query,
                new QueryStats(
                    ImmutableMap.<String, Object>of(
                        "query/time", queryTime, "success", false, "exception", e.toString()))));
      } catch (Exception e2) {
        log.error(e2, "Unable to log query [%s]!", queryString);
      }

      log.makeAlert(e, "Exception handling request")
          .addData("exception", e.toString())
          .addData("query", queryString)
          .addData("peer", req.getRemoteAddr())
          .emit();

      return Response.serverError()
          .type(contentType)
          .entity(
              jsonWriter.writeValueAsBytes(
                  ImmutableMap.of(
                      "error", e.getMessage() == null ? "null exception" : e.getMessage())))
          .build();
    } finally {
      Thread.currentThread().setName(currThreadName);
    }
  }