/*
Copyright (c) 2012 LinkedIn Corp.
Licensed 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 com.linkedin.d2.balancer;
import com.linkedin.common.callback.Callback;
import com.linkedin.common.util.None;
import com.linkedin.d2.balancer.clients.TrackerClient;
import com.linkedin.d2.balancer.clients.TrackerClientTest.TestClient;
import com.linkedin.d2.balancer.properties.ClusterProperties;
import com.linkedin.d2.balancer.properties.PartitionData;
import com.linkedin.d2.balancer.properties.ServiceProperties;
import com.linkedin.d2.balancer.properties.UriProperties;
import com.linkedin.d2.balancer.strategies.LoadBalancerStrategy;
import com.linkedin.d2.balancer.strategies.random.RandomLoadBalancerStrategy;
import com.linkedin.d2.balancer.util.partitions.DefaultPartitionAccessor;
import com.linkedin.d2.balancer.util.partitions.PartitionAccessor;
import com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback;
import com.linkedin.r2.transport.common.bridge.client.TransportClient;
import java.net.URI;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
public class LoadBalancerTestState implements LoadBalancerState
{
public boolean getClient = false;
public boolean getClusterProperties = false;
public boolean getServiceProperties = false;
public boolean getStrategy = false;
public boolean getUriProperties = false;
public boolean isListeningToCluster = false;
public boolean isListeningToService = false;
public boolean listenToService = false;
public boolean listenToCluster = false;
public boolean shutdown = false;
public boolean getPartitionAccessor = false;
@Override
public TrackerClient getClient(String clusterName, URI uri)
{
Map<Integer, PartitionData> partitionDataMap = new HashMap<Integer, PartitionData>(2);
partitionDataMap.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1));
return (getClient) ? new TrackerClient(uri, partitionDataMap, new TestClient()) : null;
}
@Override
public TransportClient getClient(String clusterName, String Scheme)
{
return (getClient) ? new TestClient() : null;
}
@Override
public LoadBalancerStateItem<ClusterProperties> getClusterProperties(String clusterName)
{
return (getClusterProperties)
? new LoadBalancerStateItem<ClusterProperties>(new ClusterProperties("cluster-1"),
0,
0) : null;
}
@Override
public LoadBalancerStateItem<PartitionAccessor> getPartitionAccessor(String clusterName)
{
return getPartitionAccessor
? new LoadBalancerStateItem<PartitionAccessor>(DefaultPartitionAccessor.getInstance(), 0, 0) : null;
}
@Override
public LoadBalancerStateItem<ServiceProperties> getServiceProperties(String serviceName)
{
List<String> prioritizedSchemes = new ArrayList<String>();
prioritizedSchemes.add("http");
return (getServiceProperties)
? new LoadBalancerStateItem<ServiceProperties>(new ServiceProperties("service-1",
"cluster-1",
"/foo",
"rr",
Collections.<String>emptyList(),
Collections.<String, Object>emptyMap(),
null,
null,
prioritizedSchemes,
null),
0,
0) : null;
}
@Override
public LoadBalancerStateItem<UriProperties> getUriProperties(String clusterName)
{
try
{
URI uri1 = URI.create("http://test.qa1.com:1234");
URI uri2 = URI.create("http://test.qa2.com:2345");
URI uri3 = URI.create("http://test.qa3.com:6789");
Map<Integer, PartitionData> partitionData = new HashMap<Integer, PartitionData>(1);
partitionData.put(DefaultPartitionAccessor.DEFAULT_PARTITION_ID, new PartitionData(1d));
Map<URI, Map<Integer, PartitionData>> uriData = new HashMap<URI, Map<Integer, PartitionData>>(3);
uriData.put(uri1, partitionData);
uriData.put(uri2, partitionData);
uriData.put(uri3, partitionData);
return (getUriProperties)
? new LoadBalancerStateItem<UriProperties>(new UriProperties("cluster-1", uriData),
0,
0) : null;
}
catch (Exception e)
{
throw new RuntimeException(e);
}
}
@Override
public LoadBalancerStrategy getStrategy(String serviceName, String scheme)
{
return (getStrategy) ? new RandomLoadBalancerStrategy() : null;
}
@Override
public boolean isListeningToCluster(String clusterName)
{
return isListeningToCluster;
}
@Override
public boolean isListeningToService(String serviceName)
{
return isListeningToService;
}
@Override
public void listenToCluster(String clusterName,
LoadBalancerStateListenerCallback callback)
{
if (listenToCluster)
{
callback.done(LoadBalancerStateListenerCallback.CLUSTER, clusterName);
}
}
@Override
public void listenToService(String serviceName,
LoadBalancerStateListenerCallback callback)
{
if (listenToService)
{
callback.done(LoadBalancerStateListenerCallback.SERVICE, serviceName);
}
}
@Override
public void start(Callback<None> callback)
{
callback.onSuccess(None.none());
}
@Override
public void shutdown(PropertyEventShutdownCallback shutdown)
{
if (this.shutdown)
{
shutdown.done();
}
}
public List<SchemeStrategyPair> getStrategiesForService(String serviceName,
List<String> prioritizedSchemes)
{
List<SchemeStrategyPair> orderedStrategies = new ArrayList<SchemeStrategyPair>(prioritizedSchemes.size());
for (String scheme : prioritizedSchemes)
{
LoadBalancerStrategy strategy = getStrategy(serviceName, scheme);
if (strategy != null)
{
orderedStrategies.add(new SchemeStrategyPair(
scheme,
getStrategy(serviceName, scheme))
);
}
}
return orderedStrategies;
}
}