Fix coordinator's dataSource api with full parameter (#5662)

* Fix coordinator's dataSource api with full parameter

* address comment

* Add a constructor for json serde and fix result order

* Change to immutableSortedMap

* Revert immutableSortedMap to treeMap

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

@@ -21,13 +21,12 @@ package io.druid.client;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableMap;
 import io.druid.timeline.DataSegment;
 
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentSkipListMap;
 
 /**
  */
@@ -35,7 +34,7 @@ public class DruidDataSource
 {
   private final String name;
   private final Map<String, String> properties;
-  private final ConcurrentHashMap<String, DataSegment> idToSegmentMap;
+  private final ConcurrentSkipListMap<String, DataSegment> idToSegmentMap;
 
   public DruidDataSource(
       String name,
@@ -44,7 +43,7 @@ public class DruidDataSource
   {
     this.name = Preconditions.checkNotNull(name);
     this.properties = properties;
-    this.idToSegmentMap = new ConcurrentHashMap<>();
+    this.idToSegmentMap = new ConcurrentSkipListMap<>();
   }
 
   @JsonProperty
@@ -83,11 +82,7 @@ public class DruidDataSource
 
   public ImmutableDruidDataSource toImmutableDruidDataSource()
   {
-    return new ImmutableDruidDataSource(
+    return new ImmutableDruidDataSource(name, properties, idToSegmentMap);
-        name,
-        ImmutableMap.copyOf(properties),
-        ImmutableMap.copyOf(idToSegmentMap)
-    );
   }
 
   @Override
@@ -102,6 +97,7 @@ public class DruidDataSource
   @Override
   public boolean equals(Object o)
   {
+    //noinspection Contract
     throw new UnsupportedOperationException("Use ImmutableDruidDataSource instead");
   }
 

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

@@ -19,12 +19,18 @@
 
 package io.druid.client;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
 import io.druid.timeline.DataSegment;
 
 import java.util.Collection;
+import java.util.Map;
 import java.util.Objects;
+import java.util.SortedMap;
 
 /**
  */
@@ -32,29 +38,52 @@ public class ImmutableDruidDataSource
 {
   private final String name;
   private final ImmutableMap<String, String> properties;
-  private final ImmutableMap<String, DataSegment> idToSegments;
+  private final ImmutableSortedMap<String, DataSegment> idToSegments;
 
   public ImmutableDruidDataSource(
       String name,
-      ImmutableMap<String, String> properties,
+      Map<String, String> properties,
-      ImmutableMap<String, DataSegment> idToSegments
+      SortedMap<String, DataSegment> idToSegments
   )
   {
     this.name = Preconditions.checkNotNull(name);
-    this.properties = properties;
+    this.properties = ImmutableMap.copyOf(properties);
-    this.idToSegments = idToSegments;
+    this.idToSegments = ImmutableSortedMap.copyOfSorted(idToSegments);
   }
 
+  @JsonCreator
+  public ImmutableDruidDataSource(
+      @JsonProperty("name") String name,
+      @JsonProperty("properties") Map<String, String> properties,
+      @JsonProperty("segments") Collection<DataSegment> segments
+  )
+  {
+    this.name = Preconditions.checkNotNull(name);
+    this.properties = ImmutableMap.copyOf(properties);
+    final ImmutableSortedMap.Builder<String, DataSegment> builder = ImmutableSortedMap.naturalOrder();
+    segments.forEach(segment -> builder.put(segment.getIdentifier(), segment));
+    this.idToSegments = builder.build();
+  }
+
+  @JsonProperty
   public String getName()
   {
     return name;
   }
 
+  @JsonProperty
+  public Map<String, String> getProperties()
+  {
+    return properties;
+  }
+
+  @JsonProperty
   public Collection<DataSegment> getSegments()
   {
     return idToSegments.values();
   }
 
+  @JsonIgnore
   public DataSegment getSegment(String segmentIdentifier)
   {
     return idToSegments.get(segmentIdentifier);

server/src/main/java/io/druid/server/http/DatasourcesResource.java

@@ -65,12 +65,15 @@ import javax.ws.rs.QueryParam;
 import javax.ws.rs.core.Context;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.SortedMap;
+import java.util.TreeMap;
 import java.util.stream.Collectors;
 
 /**
@@ -143,7 +146,7 @@ public class DatasourcesResource
       @QueryParam("full") final String full
   )
   {
-    ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
+    final ImmutableDruidDataSource dataSource = getDataSource(dataSourceName);
 
     if (dataSource == null) {
       return Response.noContent().build();
@@ -508,7 +511,7 @@ public class DatasourcesResource
       return null;
     }
 
-    Map<String, DataSegment> segmentMap = Maps.newHashMap();
+    final SortedMap<String, DataSegment> segmentMap = new TreeMap<>();
     for (ImmutableDruidDataSource dataSource : dataSources) {
       Iterable<DataSegment> segments = dataSource.getSegments();
       for (DataSegment segment : segments) {
@@ -516,11 +519,7 @@ public class DatasourcesResource
       }
     }
 
-    return new ImmutableDruidDataSource(
+    return new ImmutableDruidDataSource(dataSourceName, Collections.emptyMap(), segmentMap);
-        dataSourceName,
-        ImmutableMap.of(),
-        ImmutableMap.copyOf(segmentMap)
-    );
   }
 
   private Pair<DataSegment, Set<String>> getSegment(String segmentId)

server/src/test/java/io/druid/client/ImmutableDruidDataSourceTest.java

@@ -0,0 +1,64 @@
+/*
+ * 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.client;
+
+import com.fasterxml.jackson.databind.InjectableValues.Std;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSortedMap;
+import io.druid.jackson.DefaultObjectMapper;
+import io.druid.java.util.common.Intervals;
+import io.druid.timeline.DataSegment;
+import io.druid.timeline.DataSegment.PruneLoadSpecHolder;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.IOException;
+
+public class ImmutableDruidDataSourceTest
+{
+  @Test
+  public void testSerde() throws IOException
+  {
+    final DataSegment segment = new DataSegment(
+        "test",
+        Intervals.of("2017/2018"),
+        "version",
+        null,
+        ImmutableList.of("dim1", "dim2"),
+        ImmutableList.of("met1", "met2"),
+        null,
+        1,
+        100L,
+        PruneLoadSpecHolder.DEFAULT
+    );
+    final ImmutableDruidDataSource dataSource = new ImmutableDruidDataSource(
+        "test",
+        ImmutableMap.of("prop1", "val1", "prop2", "val2"),
+        ImmutableSortedMap.of(segment.getIdentifier(), segment)
+    );
+
+    final ObjectMapper objectMapper = new DefaultObjectMapper()
+        .setInjectableValues(new Std().addValue(PruneLoadSpecHolder.class, PruneLoadSpecHolder.DEFAULT));
+    final String json = objectMapper.writeValueAsString(dataSource);
+    Assert.assertEquals(dataSource, objectMapper.readValue(json, ImmutableDruidDataSource.class));
+  }
+}

server/src/test/java/io/druid/server/coordinator/DruidClusterTest.java

@@ -39,6 +39,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.NavigableSet;
 import java.util.Set;
+import java.util.TreeMap;
 import java.util.TreeSet;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -74,14 +75,14 @@ public class DruidClusterTest
       "src1",
       new ImmutableDruidDataSource(
           "src1",
-          ImmutableMap.of(),
+          Collections.emptyMap(),
-          ImmutableMap.of()
+          new TreeMap<>()
       ),
       "src2",
       new ImmutableDruidDataSource(
           "src2",
-          ImmutableMap.of(),
+          Collections.emptyMap(),
-          ImmutableMap.of()
+          new TreeMap<>()
       )
   );
 

server/src/test/java/io/druid/server/coordinator/ServerHolderTest.java

@@ -31,8 +31,10 @@ import io.druid.timeline.partition.NoneShardSpec;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.TreeMap;
 
 public class ServerHolderTest
 {
@@ -65,14 +67,14 @@ public class ServerHolderTest
       "src1",
       new ImmutableDruidDataSource(
           "src1",
-          ImmutableMap.of(),
+          Collections.emptyMap(),
-          ImmutableMap.of()
+          new TreeMap<>()
       ),
       "src2",
       new ImmutableDruidDataSource(
           "src2",
-          ImmutableMap.of(),
+          Collections.emptyMap(),
-          ImmutableMap.of()
+          new TreeMap<>()
       )
   );
 

server/src/test/java/io/druid/server/http/DatasourcesResourceTest.java

@@ -314,7 +314,7 @@ public class DatasourcesResourceTest
   @Test
   public void testFullGetTheDataSource()
   {
-    DruidDataSource dataSource1 = new DruidDataSource("datasource1", new HashMap());
+    DruidDataSource dataSource1 = new DruidDataSource("datasource1", new HashMap<>());
     EasyMock.expect(server.getDataSource("datasource1")).andReturn(
         dataSource1
     ).atLeastOnce();