net.snowflake.client.core.QueryContextCache Maven / Gradle / Ivy
/*
* Copyright (c) 2012-2022 Snowflake Computing Inc. All rights reserved.
*/
package net.snowflake.client.core;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.*;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
/**
* Most Recently Used and Priority based cache. A separate cache for each connection in the driver.
*/
public class QueryContextCache {
private final int capacity; // Capacity of the cache
private final HashMap idMap; // Map for id and QCC
private final TreeSet treeSet; // Order data as per priority
private final HashMap priorityMap; // Map for priority and QCC
private final HashMap
newPriorityMap; // Intermediate map for priority and QCC for current round of merging
private static final SFLogger logger = SFLoggerFactory.getLogger(QueryContextCache.class);
private static ObjectMapper jsonObjectMapper;
static {
jsonObjectMapper = new ObjectMapper();
}
/**
* Constructor.
*
* @param capacity Maximum capacity of the cache.
*/
public QueryContextCache(int capacity) {
this.capacity = capacity;
idMap = new HashMap<>();
priorityMap = new HashMap<>();
newPriorityMap = new HashMap<>();
treeSet =
new TreeSet<>(
Comparator.comparingLong(QueryContextElement::getPriority)
.thenComparingLong(QueryContextElement::getId)
.thenComparingLong(QueryContextElement::getReadTimestamp));
}
/**
* Merge a new element comes from the server with the existing cache. Merge is based on read time
* stamp for the same id and based on priority for two different ids.
*
* @param id Database id.
* @param readTimestamp Last time read metadata from FDB.
* @param priority 0 to N number, where 0 is the highest priority. Eviction policy is based on
* priority.
* @param context Opaque query context.
*/
void merge(long id, long readTimestamp, long priority, String context) {
if (idMap.containsKey(id)) {
// ID found in the cache
QueryContextElement qce = idMap.get(id);
if (readTimestamp > qce.readTimestamp) {
if (qce.priority == priority) {
// Same priority, overwrite new data at same place
qce.readTimestamp = readTimestamp;
qce.context = context;
} else {
// Change in priority
QueryContextElement newQCE =
new QueryContextElement(id, readTimestamp, priority, context);
replaceQCE(qce, newQCE);
} // new priority
} // new data is recent
else if (readTimestamp == qce.readTimestamp && qce.priority != priority) {
// Same read timestamp but change in priority
QueryContextElement newQCE = new QueryContextElement(id, readTimestamp, priority, context);
replaceQCE(qce, newQCE);
}
} // id found
else {
// new id
if (priorityMap.containsKey(priority)) {
// Same priority with different id
QueryContextElement qce = priorityMap.get(priority);
// Replace with new data
QueryContextElement newQCE = new QueryContextElement(id, readTimestamp, priority, context);
replaceQCE(qce, newQCE);
} else {
// new priority
// Add new element in the cache
QueryContextElement newQCE = new QueryContextElement(id, readTimestamp, priority, context);
addQCE(newQCE);
}
}
}
/** Sync the newPriorityMap with the priorityMap at the end of current round of merge */
void syncPriorityMap() {
logger.debug(
"syncPriorityMap called priorityMap size = {}, newPrioirtyMap size = {}",
priorityMap.size(),
newPriorityMap.size());
for (Map.Entry entry : newPriorityMap.entrySet()) {
priorityMap.put(entry.getKey(), entry.getValue());
}
// clear the newPriorityMap for next round of QCC merge(a round consists of multiple entries)
newPriorityMap.clear();
}
/**
* After the merge, loop through priority list and make sure cache is at most capacity. Remove all
* other elements from the list based on priority.
*/
void checkCacheCapacity() {
logger.debug(
"checkCacheCapacity() called. treeSet size {} cache capacity {}", treeSet.size(), capacity);
if (treeSet.size() > capacity) {
// remove elements based on priority
while (treeSet.size() > capacity) {
QueryContextElement qce = treeSet.last();
removeQCE(qce);
}
}
logger.debug(
"checkCacheCapacity() returns. treeSet size {} cache capacity {}",
treeSet.size(),
capacity);
}
/** Clear the cache. */
public void clearCache() {
logger.debug("clearCache() called");
idMap.clear();
priorityMap.clear();
treeSet.clear();
logger.debug("clearCache() returns. Number of entries in cache now {}", treeSet.size());
}
/**
* @param data: the QueryContext Object serialized as a JSON format string
*/
public void deserializeQueryContextJson(String data) {
synchronized (this) {
// Log existing cache entries
logCacheEntries();
if (data == null || data.length() == 0) {
// Clear the cache
clearCache();
return;
}
try {
JsonNode rootNode = jsonObjectMapper.readTree(data);
// Deserialize the entries. The first entry with priority is the main entry. On JDBC side,
// we save all entries into one list to simplify the logic. An example JSON is:
// {
// "entries": [
// {
// "id": 0,
// "read_timestamp": 123456789,
// "priority": 0,
// "context": "base64 encoded context"
// },
// {
// "id": 1,
// "read_timestamp": 123456789,
// "priority": 1,
// "context": "base64 encoded context"
// },
// {
// "id": 2,
// "read_timestamp": 123456789,
// "priority": 2,
// "context": "base64 encoded context"
// }
// ]
JsonNode entriesNode = rootNode.path("entries");
if (entriesNode != null && entriesNode.isArray()) {
for (JsonNode entryNode : entriesNode) {
QueryContextElement entry = deserializeQueryContextElement(entryNode);
if (entry != null) {
merge(entry.id, entry.readTimestamp, entry.priority, entry.context);
} else {
logger.warn(
"deserializeQueryContextJson: deserializeQueryContextElement meets mismatch field type. Clear the QueryContextCache.");
clearCache();
return;
}
}
// after merging all entries, sync the internal priority map to priority map. Because of
// priority swicth from GS side,
// there could be priority key conflict if we directly operating on the priorityMap during
// a round of merge.
syncPriorityMap();
}
} catch (Exception e) {
logger.debug("deserializeQueryContextJson: Exception = {}", e.getMessage());
// Not rethrowing. clear the cache as incomplete merge can lead to unexpected behavior.
clearCache();
}
// After merging all entries, truncate to capacity
checkCacheCapacity();
// Log existing cache entries
logCacheEntries();
} // Synchronized
}
private static QueryContextElement deserializeQueryContextElement(JsonNode node)
throws IOException {
QueryContextElement entry = new QueryContextElement();
JsonNode idNode = node.path("id");
if (idNode.isNumber()) {
entry.setId(idNode.asLong());
} else {
logger.warn("deserializeQueryContextElement: `id` field is not Number type");
return null;
}
JsonNode timestampNode = node.path("timestamp");
if (timestampNode.isNumber()) {
entry.setReadTimestamp(timestampNode.asLong());
} else {
logger.warn("deserializeQueryContextElement: `timestamp` field is not Long type");
return null;
}
JsonNode priorityNode = node.path("priority");
if (priorityNode.isNumber()) {
entry.setPriority(priorityNode.asLong());
} else {
logger.warn("deserializeQueryContextElement: `priority` field is not Long type");
return null;
}
JsonNode contextNode = node.path("context");
if (contextNode.isTextual()) {
String contextBytes = contextNode.asText();
entry.setContext(contextBytes);
} else if (contextNode.isEmpty()) {
// Currenly the OpaqueContext field is empty in the JSON received from GS. In the future, it
// will
// be filled with OpaqueContext object in base64 format.
logger.debug("deserializeQueryContextElement `context` field is empty");
} else {
logger.warn("deserializeQueryContextElement: `context` field is not String type");
return null;
}
return entry;
}
/**
* Deserialize the QueryContext cache from a QueryContextDTO object. This function currently is
* only used in QueryContextCacheTest.java where we check that after serialization and
* deserialization, the cache is the same as before.
*/
public void deserializeQueryContextDTO(QueryContextDTO queryContextDTO) {
synchronized (this) {
// Log existing cache entries
logCacheEntries();
if (queryContextDTO == null) {
// Clear the cache
clearCache();
// Log existing cache entries
logCacheEntries();
return;
}
try {
List entries = queryContextDTO.getEntries();
if (entries != null) {
for (QueryContextEntryDTO entryDTO : entries) {
// The main entry priority will always be 0, we simply save a list of
// QueryContextEntryDTO in QueryContextDTO
QueryContextElement entry = deserializeQueryContextElementDTO(entryDTO);
merge(entry.id, entry.readTimestamp, entry.priority, entry.context);
logCacheEntries();
}
}
// after merging all entries, sync the internal priority map to priority map. Because of
// priority swicth from GS side,
// there could be priority key conflict if we directly operating on the priorityMap during a
// round of merge.
syncPriorityMap();
} catch (Exception e) {
logger.debug("deserializeQueryContextDTO: Exception = {}", e.getMessage());
// Not rethrowing. clear the cache as incomplete merge can lead to unexpected behavior.
clearCache();
}
// After merging all entries, truncate to capacity
checkCacheCapacity();
// Log existing cache entries
logCacheEntries();
} // Synchronized
}
private static QueryContextElement deserializeQueryContextElementDTO(
QueryContextEntryDTO entryDTO) throws IOException {
QueryContextElement entry =
new QueryContextElement(
entryDTO.getId(),
entryDTO.getTimestamp(),
entryDTO.getPriority(),
entryDTO.getContext().getBase64Data());
return entry;
}
/**
* Serialize the QueryContext cache to a QueryContextDTO object, which can be serialized to JSON
* automatically later.
*/
public QueryContextDTO serializeQueryContextDTO() {
synchronized (this) {
// Log existing cache entries
logCacheEntries();
TreeSet elements = getElements();
if (elements.size() == 0) return null;
try {
QueryContextDTO queryContextDTO = new QueryContextDTO();
List entries = new ArrayList();
// the first element is the main entry with priority 0. We use a list of
// QueryContextEntryDTO to store all entries in QueryContextDTO
// to simplify the JDBC side QueryContextCache design.
for (final QueryContextElement elem : elements) {
QueryContextEntryDTO queryContextElementDTO = serializeQueryContextEntryDTO(elem);
entries.add(queryContextElementDTO);
}
queryContextDTO.setEntries(entries);
return queryContextDTO;
} catch (Exception e) {
logger.debug("serializQueryContextDTO(): Exception {}", e.getMessage());
return null;
}
}
}
private QueryContextEntryDTO serializeQueryContextEntryDTO(QueryContextElement entry)
throws IOException {
// OpaqueContextDTO contains a base64 encoded byte array. On JDBC side, we do not decode and
// encode it
QueryContextEntryDTO entryDTO =
new QueryContextEntryDTO(
entry.getId(),
entry.getReadTimestamp(),
entry.getPriority(),
new OpaqueContextDTO(entry.getContext()));
return entryDTO;
}
/**
* @param id the id of the element
* @param timestamp the last update timestamp
* @param priority the priority of the element
* @param opaqueContext the binary data of the opaque context
* @return a query context element
*/
private static QueryContextElement createElement(
long id, long timestamp, long priority, String opaqueContext) {
return new QueryContextElement(id, timestamp, priority, opaqueContext);
}
/**
* Add an element in the cache.
*
* @param qce element to add
*/
private void addQCE(QueryContextElement qce) {
idMap.put(qce.id, qce);
priorityMap.put(qce.priority, qce);
treeSet.add(qce);
}
/**
* Remove an element from the cache.
*
* @param qce element to remove.
*/
private void removeQCE(QueryContextElement qce) {
treeSet.remove(qce);
priorityMap.remove(qce.priority);
idMap.remove(qce.id);
}
/**
* Replace the cache element with a new response element. Remove old element exist in the cache
* and add a new element received.
*
* @param oldQCE an element exist in the cache
* @param newQCE a new element just received.
*/
private void replaceQCE(QueryContextElement oldQCE, QueryContextElement newQCE) {
// Remove old element from the cache
removeQCE(oldQCE);
// Add new element in the cache
addQCE(newQCE);
}
/**
* Get all elements in the cache in the order of the priority.
*
* @return TreeSet containing cache elements
*/
private TreeSet getElements() {
return treeSet;
}
int getSize() {
return treeSet.size();
}
void getElements(long[] ids, long[] readTimestamps, long[] priorities, String[] contexts) {
TreeSet elems = getElements();
int i = 0;
for (QueryContextElement elem : elems) {
ids[i] = elem.id;
readTimestamps[i] = elem.readTimestamp;
priorities[i] = elem.priority;
contexts[i] = elem.context;
i++;
}
}
/** Debugging purpose, log the all entries in the cache. */
void logCacheEntries() {
if (logger.isDebugEnabled()) {
TreeSet elements = getElements();
for (final QueryContextElement elem : elements) {
logger.debug(
" Cache Entry: id: {} readTimestamp: {} priority: {}",
elem.id,
elem.readTimestamp,
elem.priority);
}
}
}
/** Query context information. */
private static class QueryContextElement implements Comparable {
long id; // database id as key. (bigint)
long readTimestamp; // When the query context read (bigint). Compare for same id.
long priority; // Priority of the query context (bigint). Compare for different ids.
String context; // Opaque information (varbinary).
public QueryContextElement() {
// Default constructor
}
/**
* Constructor.
*
* @param id database id
* @param readTimestamp Server time when this entry read
* @param priority Priority of this entry w.r.t other ids
* @param context Opaque query context, used by query processor in the server.
*/
public QueryContextElement(long id, long readTimestamp, long priority, String context) {
this.id = id;
this.readTimestamp = readTimestamp;
this.priority = priority;
this.context = context;
}
@Override
public boolean equals(Object obj) {
if (obj == this) {
return true;
}
if (!(obj instanceof QueryContextElement)) {
return super.equals(obj);
}
QueryContextElement other = (QueryContextElement) obj;
return (id == other.id
&& readTimestamp == other.readTimestamp
&& priority == other.priority
&& context.equals(other.context));
}
@Override
public int hashCode() {
int hash = 31;
hash = hash * 31 + (int) id;
hash += (hash * 31) + (int) readTimestamp;
hash += (hash * 31) + (int) priority;
hash += (hash * 31) + context.hashCode();
return hash;
}
/**
* Keep elements in ascending order of the priority. This method called by TreeSet.
*
* @param obj the object to be compared.
* @return 0 if equals, -1 if this element is less than new element, otherwise 1.
*/
public int compareTo(QueryContextElement obj) {
return (priority == obj.priority) ? 0 : (((priority - obj.priority) < 0) ? -1 : 1);
}
public void setId(long id) {
this.id = id;
}
public void setPriority(long priority) {
this.priority = priority;
}
public void setContext(String context) {
this.context = context;
}
public void setReadTimestamp(long readTimestamp) {
this.readTimestamp = readTimestamp;
}
public long getId() {
return id;
}
public long getReadTimestamp() {
return readTimestamp;
}
public long getPriority() {
return priority;
}
public String getContext() {
return context;
}
}
}
© 2015 - 2024 Weber Informatics LLC | Privacy Policy