Make realtimes available for loading segments

server/src/main/java/io/druid/client/CachingClusteredClient.java

@@ -272,7 +272,7 @@ public ShardSpec apply(PartitionChunk<ServerSelector> input)
       Hasher hasher = Hashing.sha1().newHasher();
       boolean hasOnlyHistoricalSegments = true;
       for (Pair<ServerSelector, SegmentDescriptor> p : segments) {
-        if (!p.lhs.pick().getServer().isAssignable()) {
+        if (!p.lhs.pick().getServer().segmentReplicatable()) {
           hasOnlyHistoricalSegments = false;
           break;
         }
@@ -429,7 +429,7 @@ private void addSequencesFromServer(ArrayList<Sequence<T>> listOfSequences)
               final MultipleSpecificSegmentSpec segmentSpec = new MultipleSpecificSegmentSpec(descriptors);
 
               final Sequence<T> resultSeqToAdd;
-              if (!server.isAssignable() || !populateCache || isBySegment) { // Direct server queryable
+              if (!server.segmentReplicatable() || !populateCache || isBySegment) { // Direct server queryable
                 if (!isBySegment) {
                   resultSeqToAdd = clientQueryable.run(queryPlus.withQuerySegmentSpec(segmentSpec), responseContext);
                 } else {

server/src/main/java/io/druid/client/DruidServer.java

@@ -28,6 +28,7 @@
 import io.druid.java.util.common.logger.Logger;
 import io.druid.server.DruidNode;
 import io.druid.server.coordination.DruidServerMetadata;
+import io.druid.server.coordination.ServerType;
 import io.druid.timeline.DataSegment;
 
 import java.util.Collections;
@@ -111,7 +112,7 @@ public long getMaxSize()
     return metadata.getMaxSize();
   }
 
-  public String getType()
+  public ServerType getType()
   {
     return metadata.getType();
   }
@@ -121,9 +122,9 @@ public String getTier()
     return metadata.getTier();
   }
 
-  public boolean isAssignable()
+  public boolean segmentReplicatable()
   {
-    return metadata.isAssignable();
+    return metadata.segmentReplicatable();
   }
 
   public int getPriority()

server/src/main/java/io/druid/client/ImmutableDruidDataSource.java

@@ -72,4 +72,15 @@ public Set<DataSegment> getSegments()
   {
     return segmentsHolder;
   }
+
+  @Override
+  public String toString()
+  {
+    // partitionNames is intentionally ignored because it is usually large
+    return "ImmutableDruidDataSource{"
+           + "name='" + name
+           + "', segments='" + segmentsHolder
+           + "', properties='" + properties
+           + "'}";
+  }
 }

server/src/main/java/io/druid/client/ImmutableDruidServer.java

@@ -21,6 +21,7 @@
 
 import com.google.common.collect.ImmutableMap;
 import io.druid.server.coordination.DruidServerMetadata;
+import io.druid.server.coordination.ServerType;
 import io.druid.timeline.DataSegment;
 
 import java.util.Map;
@@ -72,7 +73,7 @@ public long getMaxSize()
     return metadata.getMaxSize();
   }
 
-  public String getType()
+  public ServerType getType()
   {
     return metadata.getType();
   }
@@ -106,4 +107,15 @@ public Map<String, DataSegment> getSegments()
   {
     return segments;
   }
+
+  @Override
+  public String toString()
+  {
+    // segments is intentionally ignored because it is usually large
+    return "ImmutableDruidServer{"
+           + "meta='" + metadata
+           + "', size='" + currSize
+           + "', sources='" + dataSources
+           + "'}";
+  }
 }

server/src/main/java/io/druid/segment/realtime/RealtimeManager.java

@@ -20,6 +20,7 @@
 package io.druid.segment.realtime;
 
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Function;
 import com.google.common.base.Supplier;
 import com.google.common.base.Throwables;
@@ -84,6 +85,7 @@ public RealtimeManager(
     this(fireDepartments, conglomerate, serverAnnouncer, Maps.newHashMap());
   }
 
+  @VisibleForTesting
   RealtimeManager(
       List<FireDepartment> fireDepartments,
       QueryRunnerFactoryConglomerate conglomerate,
@@ -94,7 +96,7 @@ public RealtimeManager(
     this.fireDepartments = fireDepartments;
     this.conglomerate = conglomerate;
     this.serverAnnouncer = serverAnnouncer;
-    this.chiefs = chiefs;
+    this.chiefs = chiefs == null ? Maps.newHashMap() : Maps.newHashMap(chiefs);
   }
 
   @LifecycleStart

server/src/main/java/io/druid/segment/realtime/plumber/CoordinatorBasedSegmentHandoffNotifier.java

@@ -148,7 +148,7 @@ static boolean isHandOffComplete(List<ImmutableSegmentLoadInfo> serverView, Segm
             @Override
             public boolean apply(DruidServerMetadata input)
             {
-              return input.isAssignable();
+              return input.segmentReplicatable();
             }
           }
       )) {

server/src/main/java/io/druid/server/SegmentManager.java

@@ -0,0 +1,196 @@
+/*
+ * Licensed to Metamarkets Group Inc. (Metamarkets) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. Metamarkets licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.druid.server;
+
+import com.google.common.collect.Ordering;
+import com.google.inject.Inject;
+import com.metamx.emitter.EmittingLogger;
+import io.druid.collections.CountingMap;
+import io.druid.segment.ReferenceCountingSegment;
+import io.druid.segment.Segment;
+import io.druid.segment.loading.SegmentLoader;
+import io.druid.segment.loading.SegmentLoadingException;
+import io.druid.timeline.DataSegment;
+import io.druid.timeline.VersionedIntervalTimeline;
+import io.druid.timeline.partition.PartitionChunk;
+import io.druid.timeline.partition.PartitionHolder;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+
+public class SegmentManager
+{
+  private static final EmittingLogger log = new EmittingLogger(SegmentManager.class);
+
+  private final Object lock = new Object();
+  private final SegmentLoader segmentLoader;
+  private final Map<String, VersionedIntervalTimeline<String, ReferenceCountingSegment>> dataSources = new HashMap<>();
+  private final CountingMap<String> dataSourceSizes = new CountingMap<>();
+  private final CountingMap<String> dataSourceCounts = new CountingMap<>();
+
+  @Inject
+  public SegmentManager(
+      SegmentLoader segmentLoader
+  )
+  {
+    this.segmentLoader = segmentLoader;
+  }
+
+  public Map<String, Long> getDataSourceSizes()
+  {
+    synchronized (dataSourceSizes) {
+      return dataSourceSizes.snapshot();
+    }
+  }
+
+  public Map<String, Long> getDataSourceCounts()
+  {
+    synchronized (dataSourceCounts) {
+      return dataSourceCounts.snapshot();
+    }
+  }
+
+  public boolean isSegmentCached(final DataSegment segment) throws SegmentLoadingException
+  {
+    return segmentLoader.isSegmentLoaded(segment);
+  }
+
+  public VersionedIntervalTimeline<String, ReferenceCountingSegment> getTimeline(String dataSource)
+  {
+    synchronized (lock) {
+      return dataSources.get(dataSource);
+    }
+  }
+
+  /**
+   * Load a single segment.
+   *
+   * @param segment segment to load
+   *
+   * @return true if the segment was newly loaded, false if it was already loaded
+   *
+   * @throws SegmentLoadingException if the segment cannot be loaded
+   */
+  public boolean loadSegment(final DataSegment segment) throws SegmentLoadingException
+  {
+    final Segment adapter = getAdapter(segment);
+
+    synchronized (lock) {
+      final String dataSource = segment.getDataSource();
+      final VersionedIntervalTimeline<String, ReferenceCountingSegment> loadedIntervals = dataSources.computeIfAbsent(
+          dataSource,
+          k -> new VersionedIntervalTimeline<>(Ordering.natural())
+      );
+
+      final PartitionHolder<ReferenceCountingSegment> entry = loadedIntervals.findEntry(
+          segment.getInterval(),
+          segment.getVersion()
+      );
+      if ((entry != null) && (entry.getChunk(segment.getShardSpec().getPartitionNum()) != null)) {
+        log.warn("Told to load a adapter for a segment[%s] that already exists", segment.getIdentifier());
+        return false;
+      }
+
+      loadedIntervals.add(
+          segment.getInterval(),
+          segment.getVersion(),
+          segment.getShardSpec().createChunk(new ReferenceCountingSegment(adapter))
+      );
+      synchronized (dataSourceSizes) {
+        dataSourceSizes.add(dataSource, segment.getSize());
+      }
+      synchronized (dataSourceCounts) {
+        dataSourceCounts.add(dataSource, 1L);
+      }
+      return true;
+    }
+  }
+
+  private Segment getAdapter(final DataSegment segment) throws SegmentLoadingException
+  {
+    final Segment adapter;
+    try {
+      adapter = segmentLoader.getSegment(segment);
+    }
+    catch (SegmentLoadingException e) {
+      try {
+        segmentLoader.cleanup(segment);
+      }
+      catch (SegmentLoadingException e1) {
+        e.addSuppressed(e1);
+      }
+      throw e;
+    }
+
+    if (adapter == null) {
+      throw new SegmentLoadingException("Null adapter from loadSpec[%s]", segment.getLoadSpec());
+    }
+    return adapter;
+  }
+
+  public void dropSegment(final DataSegment segment) throws SegmentLoadingException
+  {
+    String dataSource = segment.getDataSource();
+    synchronized (lock) {
+      VersionedIntervalTimeline<String, ReferenceCountingSegment> loadedIntervals = dataSources.get(dataSource);
+
+      if (loadedIntervals == null) {
+        log.info("Told to delete a queryable for a dataSource[%s] that doesn't exist.", dataSource);
+        return;
+      }
+
+      PartitionChunk<ReferenceCountingSegment> removed = loadedIntervals.remove(
+          segment.getInterval(),
+          segment.getVersion(),
+          segment.getShardSpec().createChunk(null)
+      );
+      ReferenceCountingSegment oldQueryable = (removed == null) ? null : removed.getObject();
+
+      if (oldQueryable != null) {
+        synchronized (dataSourceSizes) {
+          dataSourceSizes.add(dataSource, -segment.getSize());
+        }
+        synchronized (dataSourceCounts) {
+          dataSourceCounts.add(dataSource, -1L);
+        }
+
+        try {
+          log.info("Attempting to close segment %s", segment.getIdentifier());
+          oldQueryable.close();
+        }
+        catch (IOException e) {
+          log.makeAlert(e, "Exception closing segment")
+             .addData("dataSource", dataSource)
+             .addData("segmentId", segment.getIdentifier())
+             .emit();
+        }
+      } else {
+        log.info(
+            "Told to delete a queryable on dataSource[%s] for interval[%s] and version [%s] that I don't have.",
+            dataSource,
+            segment.getInterval(),
+            segment.getVersion()
+        );
+      }
+    }
+    segmentLoader.cleanup(segment);
+  }
+}

server/src/main/java/io/druid/server/coordination/BatchDataSegmentAnnouncer.java

@@ -310,7 +310,7 @@ private String makeServedSegmentPath()
     return makeServedSegmentPath(
         UUIDUtils.generateUuid(
             server.getHost(),
-            server.getType(),
+            server.getType().toString(),
             server.getTier(),
             new DateTime().toString()
         )

server/src/main/java/io/druid/server/coordination/DruidServerMetadata.java

@@ -30,7 +30,7 @@ public class DruidServerMetadata
   private final String host;
   private final long maxSize;
   private final String tier;
-  private final String type;
+  private final ServerType type;
   private final int priority;
 
   @JsonCreator
@@ -47,7 +47,7 @@ public DruidServerMetadata(
     this.host = host;
     this.maxSize = maxSize;
     this.tier = tier;
-    this.type = type;
+    this.type = ServerType.fromString(type);
     this.priority = priority;
   }
 
@@ -76,7 +76,7 @@ public String getTier()
   }
 
   @JsonProperty
-  public String getType()
+  public ServerType getType()
   {
     return type;
   }
@@ -87,9 +87,9 @@ public int getPriority()
     return priority;
   }
 
-  public boolean isAssignable()
+  public boolean segmentReplicatable()
   {
-    return getType().equalsIgnoreCase("historical") || getType().equalsIgnoreCase("bridge");
+    return type.isSegmentReplicationTarget();
   }
 
   @Override