Merge pull request #364 from cliveseldon/engine_cache_fix

Fix URICache bug in engine

engine/src/main/java/io/seldon/engine/service/InternalPredictionService.java

@@ -101,7 +101,7 @@ public class InternalPredictionService {
     private final GrpcChannelHandler grpcChannelHandler;
 
     private final Map<String,HttpHeaders> headersCache = new ConcurrentHashMap<>();
-    private final Map<Endpoint,URI> uriCache = new ConcurrentHashMap<>();
+    private final Map<String,URI> uriCache = new ConcurrentHashMap<>();
 
     @Autowired
     public InternalPredictionService(RestTemplateBuilder restTemplateBuilder,AnnotationsConfig annotations,GrpcChannelHandler grpcChannelHandler,TracingProvider tracingProvider){
@@ -344,23 +344,30 @@ private boolean isDefaultData(SeldonMessage message){
 			return true;
     	return false;
     }
+
+    public static String getUriKey(Endpoint endpoint,String path)
+    {
+    	StringBuilder sb = new StringBuilder();
+    	return sb.append(endpoint.getServiceHost()).append(":").append(endpoint.getServicePort()).append(path).toString();
+    }
 
 	private SeldonMessage queryREST(String path, String dataString, PredictiveUnitState state, Endpoint endpoint, boolean isDefault)
 	{
 		long timeNow = System.currentTimeMillis();
 		URI uri;
 		try 
 		{
-			if (uriCache.containsKey(endpoint))
-				uri = uriCache.get(endpoint);
+			final String uriKey = getUriKey(endpoint, path);
+			if (uriCache.containsKey(uriKey))
+				uri = uriCache.get(uriKey);
 			else
 			{
 				URIBuilder builder = new URIBuilder().setScheme("http")
 						.setHost(endpoint.getServiceHost())
 						.setPort(endpoint.getServicePort())
 						.setPath("/"+path);
 				uri = builder.build();
-				uriCache.put(endpoint, uri);
+				uriCache.put(uriKey, uri);
 			}
 		} catch (URISyntaxException e) 
 		{

engine/src/test/java/io/seldon/engine/service/UriCacheTest.java

@@ -0,0 +1,36 @@
+package io.seldon.engine.service;
+
+import org.junit.Test;
+import org.junit.Assert;
+import io.seldon.protos.DeploymentProtos.Endpoint;
+
+
+public class UriCacheTest {
+
+	@Test
+	public void testUri()
+	{
+		Endpoint endpointA = Endpoint.newBuilder().setServiceHost("hostA").setServicePort(1000).build();
+		Endpoint endpointA2 = Endpoint.newBuilder().setServiceHost("hostA").setServicePort(1000).build();
+		Endpoint endpointB = Endpoint.newBuilder().setServiceHost("hostB").setServicePort(1000).build();
+		final String predictPath = "/predict";
+		final String predictPath2 = "/predict";
+		final String feedbackPath = "/feedback";
+
+		final String key1 = InternalPredictionService.getUriKey(endpointA, predictPath);
+		final String key2 = InternalPredictionService.getUriKey(endpointB, predictPath);
+		final String key3 = InternalPredictionService.getUriKey(endpointA, feedbackPath);
+
+		Assert.assertNotEquals(key1, key2);
+		Assert.assertNotEquals(key1, key3);
+
+		final String key4 = InternalPredictionService.getUriKey(endpointA2, predictPath);
+
+		Assert.assertEquals(key1, key4);
+
+		final String key5 = InternalPredictionService.getUriKey(endpointA, predictPath2);
+
+		Assert.assertEquals(key1, key5);
+	}
+
+}