package com.jivesoftware.os.amza.client.collection;
import com.google.common.collect.Maps;
import com.jivesoftware.os.amza.api.PartitionClient;
import com.jivesoftware.os.amza.api.partition.Consistency;
import com.jivesoftware.os.amza.api.partition.PartitionName;
import com.jivesoftware.os.amza.api.partition.PartitionProperties;
import com.jivesoftware.os.amza.client.http.AmzaClientProvider;
import com.jivesoftware.os.mlogger.core.MetricLogger;
import com.jivesoftware.os.mlogger.core.MetricLoggerFactory;
import com.jivesoftware.os.routing.bird.http.client.HttpClient;
import com.jivesoftware.os.routing.bird.shared.HttpClientException;
import java.nio.charset.StandardCharsets;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
/**
* Created by jonathan.colt on 1/16/17.
*/
public class AmzaMap<K, V> {
private static final MetricLogger LOG = MetricLoggerFactory.getLogger();
private final AmzaClientProvider<HttpClient, HttpClientException> clientProvider;
private final PartitionProperties partitionProperties;
private final int ringSize = 3; //TODO expose to conf?
private final String partitionName;
private final AmzaMarshaller<K> keyMarshaller;
private final AmzaMarshaller<V> valueMarshaller;
private final long additionalSolverAfterNMillis = 1_000; //TODO expose to conf?
private final long abandonLeaderSolutionAfterNMillis = 5_000; //TODO expose to conf?
private final long abandonSolutionAfterNMillis = 30_000; //TODO expose to conf?
public AmzaMap(AmzaClientProvider clientProvider,
String partitionName,
PartitionProperties partitionProperties,
AmzaMarshaller<K> keyMarshaller,
AmzaMarshaller<V> valueMarshaller) {
this.clientProvider = clientProvider;
this.partitionProperties = partitionProperties;
this.partitionName = partitionName;
this.keyMarshaller = keyMarshaller;
this.valueMarshaller = valueMarshaller;
}
private PartitionName partitionName() {
byte[] nameAsBytes = partitionName.getBytes(StandardCharsets.UTF_8);
return new PartitionName(false, nameAsBytes, nameAsBytes);
}
public void multiPutIfAbsent(Map<K, V> puts) throws Exception {
Map<K, V> got = multiGet(puts.keySet());
Map<K, V> put = Maps.newHashMap();
for (Entry<K, V> entry : puts.entrySet()) {
if (!got.containsKey(entry.getKey())) {
put.put(entry.getKey(), entry.getValue());
}
}
if (!put.isEmpty()) {
multiPut(put);
}
}
public void multiPut(Map<K, V> puts) throws Exception {
PartitionClient partition = clientProvider.getPartition(partitionName(), ringSize, partitionProperties);
long now = System.currentTimeMillis();
partition.commit(Consistency.leader_quorum,
null,
(stream) -> {
for (Entry<K, V> e : puts.entrySet()) {
stream.commit(keyMarshaller.toBytes(e.getKey()), valueMarshaller.toBytes(e.getValue()), now, false);
}
return true;
},
additionalSolverAfterNMillis,
abandonSolutionAfterNMillis,
Optional.empty());
LOG.info("Put ({}) {} configs.", puts.size(), partitionName);
}
public void multiRemove(List<K> keys) throws Exception {
PartitionClient partition = clientProvider.getPartition(partitionName(), ringSize, partitionProperties);
long now = System.currentTimeMillis();
partition.commit(Consistency.leader_quorum,
null,
(stream) -> {
for (K key : keys) {
stream.commit(keyMarshaller.toBytes(key), null, now, true);
}
return true;
},
additionalSolverAfterNMillis,
abandonSolutionAfterNMillis,
Optional.empty());
LOG.info("Removed ({}) {} configs.", keys.size(), partitionName);
}
public V get(K key) throws Exception {
Object[] got = new Object[1];
PartitionClient partition = clientProvider.getPartition(partitionName(), ringSize, partitionProperties);
partition.get(Consistency.leader_quorum,
null,
(keyStream) -> keyStream.stream(keyMarshaller.toBytes(key)),
(prefix, key1, value, timestamp, version) -> {
if (value != null) {
got[0] = valueMarshaller.fromBytes(value);
}
return true;
},
additionalSolverAfterNMillis,
abandonLeaderSolutionAfterNMillis,
abandonSolutionAfterNMillis,
Optional.empty()
);
LOG.info("Got {} config.",partitionName);
return (V) got[0];
}
public Map<K, V> multiGet(Collection<K> keys) throws Exception {
if (keys.isEmpty()) {
return Collections.emptyMap();
}
Map<K, V> got = Maps.newConcurrentMap();
PartitionClient partition = clientProvider.getPartition(partitionName(), ringSize, partitionProperties);
partition.get(Consistency.leader_quorum,
null,
(keyStream) -> {
for (K key : keys) {
if (!keyStream.stream(keyMarshaller.toBytes(key))) {
return false;
}
}
return true;
},
(prefix, key, value, timestamp, version) -> {
if (value != null) {
got.put(keyMarshaller.fromBytes(key), valueMarshaller.fromBytes(value));
}
return true;
},
additionalSolverAfterNMillis,
abandonLeaderSolutionAfterNMillis,
abandonSolutionAfterNMillis,
Optional.empty()
);
LOG.info("Got ({}) {} configs.", got.size(), partitionName);
return got;
}
public Map<K, V> getAll() throws Exception {
Map<K, V> got = Maps.newConcurrentMap();
PartitionClient partition = clientProvider.getPartition(partitionName(), ringSize, partitionProperties);
partition.scan(Consistency.leader_quorum,
false,
stream -> stream.stream(null, null, null, null),
(prefix, key, value, timestamp, version) -> {
if (value != null) {
got.put(keyMarshaller.fromBytes(key), valueMarshaller.fromBytes(value));
}
return true;
},
additionalSolverAfterNMillis,
abandonLeaderSolutionAfterNMillis,
abandonSolutionAfterNMillis,
Optional.empty()
);
LOG.info("Got All ({}) {} configs.", got.size(), partitionName);
return got;
}
}