src/main/java/com/googlesource/gerrit/modules/cache/pg/PgSqlStore.java - modules/cache-postgres - Git at Google

 // Copyright (C) 2018 The Android Open Source Project
 //
 // 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.googlesource.gerrit.modules.cache.pg;

 import com.google.common.cache.Cache;
 import com.google.common.hash.BloomFilter;
 import com.google.gerrit.common.TimeUtil;
 import com.google.gerrit.server.cache.PersistentCache.DiskStats;
 import com.google.inject.TypeLiteral;
 import java.io.InvalidClassException;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.sql.Timestamp;
 import java.sql.Types;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.atomic.AtomicLong;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public class PgSqlStore<K, V> {
   static final Logger log = LoggerFactory.getLogger(PgSqlStore.class);

   private final PgSqlSource source;
   private final String name;
   private final KeyType<K> keyType;
   private final long maxSize;
   private final long expireAfterWrite;
   private final BlockingQueue<PgSqlHandle> handles;
   private final AtomicLong hitCount = new AtomicLong();
   private final AtomicLong missCount = new AtomicLong();
   private volatile BloomFilter<K> bloomFilter;
   private int estimatedSize;

   private final String qCount;
   private final String qKeys;
   private final String qValue;
   private final String qTouch;
   private final String qPut;
   private final String qInvalidateKey;
   private final String qInvalidateAll;
   private final String qSum;
   private final String qOrderByAccessed;
   private final String qStats;

   PgSqlStore(
       PgSqlSource source,
       String name,
       TypeLiteral<K> keyType,
       long maxSize,
       long expireAfterWrite) {
     this.source = source;
     this.name = name;
     this.keyType = KeyType.create(keyType);
     this.maxSize = maxSize;
     this.expireAfterWrite = expireAfterWrite;

     int cores = Runtime.getRuntime().availableProcessors();
     int keep = Math.min(cores, 16);
     this.handles = new ArrayBlockingQueue<>(keep);

     // initiate all query strings
     this.qCount = "SELECT COUNT(*) FROM \"data_" + this.name + "\"";
     this.qKeys = "SELECT k FROM \"data_" + this.name + "\"";
     this.qValue = "SELECT v, created FROM \"data_" + this.name + "\" WHERE k=?";
     this.qTouch = "UPDATE \"data_" + this.name + "\" SET accessed=? WHERE k=?";
     this.qPut =
         "INSERT INTO \"data_"
             + this.name
             + "\" (k, v, created, accessed) VALUES(?,?,?,?) "
             + "ON CONFLICT (k) DO UPDATE "
             + "SET v = EXCLUDED.v, created = EXCLUDED.created, accessed = EXCLUDED.accessed";
     this.qInvalidateKey = "DELETE FROM \"data_" + this.name + "\" WHERE k=?";
     this.qInvalidateAll = "DELETE FROM \"data_" + this.name + "\"";
     this.qSum = "SELECT SUM(space) FROM \"data_" + this.name + "\"";
     this.qOrderByAccessed =
         "SELECT k,space,created FROM \"data_" + this.name + "\" ORDER BY accessed";
     this.qStats = "SELECT COUNT(*),SUM(space) FROM \"data_" + this.name + "\"";
   }

   synchronized void open() {
     if (bloomFilter == null) {
       bloomFilter = buildBloomFilter();
     }
   }

   void close() {
     PgSqlHandle h;
     while ((h = handles.poll()) != null) {
       h.close();
     }
   }

   boolean mightContain(K key) {
     BloomFilter<K> b = bloomFilter;
     if (b == null) {
       synchronized (this) {
         b = bloomFilter;
         if (b == null) {
           b = buildBloomFilter();
           bloomFilter = b;
         }
       }
     }
     return b == null || b.mightContain(key);
   }

   private BloomFilter<K> buildBloomFilter() {
     PgSqlHandle c = null;
     try {
       c = acquire();
       try (Statement s = c.conn.createStatement()) {
         if (estimatedSize <= 0) {
           try (ResultSet r = s.executeQuery(qCount)) {
             estimatedSize = r.next() ? r.getInt(1) : 0;
           }
         }

         BloomFilter<K> b = newBloomFilter();
         try (ResultSet r = s.executeQuery(qKeys)) {
           while (r.next()) {
             b.put(keyType.get(r, 1));
           }
         } catch (SQLException e) {
           if (e.getCause() instanceof InvalidClassException) {
             log.warn(
                 "Entries cached for {} "
                     + "have an incompatible class and can't be deserialized. "
                     + "Cache is flushed.",
                 name);
             invalidateAll();
           } else {
             throw e;
           }
         }
         return b;
       }
     } catch (SQLException e) {
       log.warn("Cannot build BloomFilter for {}: {}", name, e.getMessage());
       c = close(c);
       return null;
     } finally {
       release(c);
     }
   }

   ValueHolder<V> getIfPresent(K key) {
     PgSqlHandle c = null;
     try {
       c = acquire();
       if (c.get == null) {
         c.get = c.conn.prepareStatement(qValue);
       }
       keyType.set(c.get, 1, key);
       try (ResultSet r = c.get.executeQuery()) {
         if (!r.next()) {
           missCount.incrementAndGet();
           return null;
         }

         Timestamp created = r.getTimestamp(2);
         if (expired(created)) {
           invalidate(key);
           missCount.incrementAndGet();
           return null;
         }

         @SuppressWarnings("unchecked")
         V val = (V) PgSqlHandle.deserialize(r.getBytes(1));
         ValueHolder<V> h = new ValueHolder<>(val);
         h.clean = true;
         hitCount.incrementAndGet();
         touch(c, key);
         return h;
       } finally {
         c.get.clearParameters();
       }
     } catch (SQLException e) {
       log.warn("Cannot read cache {} for {}", name, key, e);
       c = close(c);
       return null;
     } finally {
       release(c);
     }
   }

   private boolean expired(Timestamp created) {
     if (expireAfterWrite == 0) {
       return false;
     }
     long age = TimeUtil.nowMs() - created.getTime();
     return 1000 * expireAfterWrite < age;
   }

   private void touch(PgSqlHandle c, K key) throws SQLException {
     if (c.touch == null) {
       c.touch = c.conn.prepareStatement(qTouch);
     }
     try {
       c.touch.setTimestamp(1, TimeUtil.nowTs());
       keyType.set(c.touch, 2, key);
       c.touch.executeUpdate();
     } finally {
       c.touch.clearParameters();
     }
   }

   void put(K key, ValueHolder<V> holder) {
     if (holder.clean) {
       return;
     }

     BloomFilter<K> b = bloomFilter;
     if (b != null) {
       b.put(key);
       bloomFilter = b;
     }

     PgSqlHandle c = null;
     try {
       c = acquire();
       if (c.put == null) {
         c.put = c.conn.prepareStatement(qPut);
       }
       try {
         keyType.set(c.put, 1, key);
         c.put.setObject(2, PgSqlHandle.serialize(holder.value), Types.BINARY);
         c.put.setTimestamp(3, new Timestamp(holder.created));
         c.put.setTimestamp(4, TimeUtil.nowTs());
         c.put.executeUpdate();
         holder.clean = true;
       } finally {
         c.put.clearParameters();
       }
     } catch (SQLException e) {
       log.warn("Cannot put into cache {}", name, e);
       c = close(c);
     } finally {
       release(c);
     }
   }

   void invalidate(K key) {
     PgSqlHandle c = null;
     try {
       c = acquire();
       invalidate(c, key);
     } catch (SQLException e) {
       log.warn("Cannot invalidate cache {}", name, e);
       c = close(c);
     } finally {
       release(c);
     }
   }

   private void invalidate(PgSqlHandle c, K key) throws SQLException {
     if (c.invalidate == null) {
       c.invalidate = c.conn.prepareStatement(qInvalidateKey);
     }
     try {
       keyType.set(c.invalidate, 1, key);
       c.invalidate.executeUpdate();
     } finally {
       c.invalidate.clearParameters();
     }
   }

   void invalidateAll() {
     PgSqlHandle c = null;
     try {
       c = acquire();
       try (Statement s = c.conn.createStatement()) {
         s.executeUpdate(qInvalidateAll);
       }
       bloomFilter = newBloomFilter();
     } catch (SQLException e) {
       log.warn("Cannot invalidate cache {}", name, e);
       c = close(c);
     } finally {
       release(c);
     }
   }

   void prune(Cache<K, ?> mem) {
     PgSqlHandle c = null;
     try {
       c = acquire();
       try (Statement s = c.conn.createStatement()) {
         long used = 0;
         try (ResultSet r = s.executeQuery(qSum)) {
           used = r.next() ? r.getLong(1) : 0;
         }
         if (used <= maxSize) {
           return;
         }

         try (ResultSet r = s.executeQuery(qOrderByAccessed)) {
           while (maxSize < used && r.next()) {
             K key = keyType.get(r, 1);
             Timestamp created = r.getTimestamp(3);
             if (mem.getIfPresent(key) != null && !expired(created)) {
               touch(c, key);
             } else {
               invalidate(c, key);
               used -= r.getLong(2);
             }
           }
         }
       }
     } catch (SQLException e) {
       log.warn("Cannot prune cache {}", name, e);
       c = close(c);
     } finally {
       release(c);
     }
   }

   DiskStats diskStats() {
     long size = 0;
     long space = 0;
     PgSqlHandle c = null;
     try {
       c = acquire();
       try (Statement s = c.conn.createStatement();
           ResultSet r = s.executeQuery(qStats)) {
         if (r.next()) {
           size = r.getLong(1);
           space = r.getLong(2);
         }
       }
     } catch (SQLException e) {
       log.warn("Cannot get DiskStats for {}", name, e);
       c = close(c);
     } finally {
       release(c);
     }
     return new DiskStats(size, space, hitCount.get(), missCount.get());
   }

   private PgSqlHandle acquire() throws SQLException {
     PgSqlHandle h = handles.poll();
     return h != null ? h : new PgSqlHandle(source, name, keyType);
   }

   private void release(PgSqlHandle h) {
     if (h != null && !handles.offer(h)) {
       h.close();
     }
   }

   private PgSqlHandle close(PgSqlHandle h) {
     if (h != null) {
       h.close();
     }
     return null;
   }

   private BloomFilter<K> newBloomFilter() {
     int cnt = Math.max(64 * 1024, 2 * estimatedSize);
     return BloomFilter.create(keyType.funnel(), cnt);
   }
 }
	// Copyright (C) 2018 The Android Open Source Project
	//
	// 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.googlesource.gerrit.modules.cache.pg;

	import com.google.common.cache.Cache;
	import com.google.common.hash.BloomFilter;
	import com.google.gerrit.common.TimeUtil;
	import com.google.gerrit.server.cache.PersistentCache.DiskStats;
	import com.google.inject.TypeLiteral;
	import java.io.InvalidClassException;
	import java.sql.ResultSet;
	import java.sql.SQLException;
	import java.sql.Statement;
	import java.sql.Timestamp;
	import java.sql.Types;
	import java.util.concurrent.ArrayBlockingQueue;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.atomic.AtomicLong;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public class PgSqlStore<K, V> {
	static final Logger log = LoggerFactory.getLogger(PgSqlStore.class);

	private final PgSqlSource source;
	private final String name;
	private final KeyType<K> keyType;
	private final long maxSize;
	private final long expireAfterWrite;
	private final BlockingQueue<PgSqlHandle> handles;
	private final AtomicLong hitCount = new AtomicLong();
	private final AtomicLong missCount = new AtomicLong();
	private volatile BloomFilter<K> bloomFilter;
	private int estimatedSize;

	private final String qCount;
	private final String qKeys;
	private final String qValue;
	private final String qTouch;
	private final String qPut;
	private final String qInvalidateKey;
	private final String qInvalidateAll;
	private final String qSum;
	private final String qOrderByAccessed;
	private final String qStats;

	PgSqlStore(
	PgSqlSource source,
	String name,
	TypeLiteral<K> keyType,
	long maxSize,
	long expireAfterWrite) {
	this.source = source;
	this.name = name;
	this.keyType = KeyType.create(keyType);
	this.maxSize = maxSize;
	this.expireAfterWrite = expireAfterWrite;

	int cores = Runtime.getRuntime().availableProcessors();
	int keep = Math.min(cores, 16);
	this.handles = new ArrayBlockingQueue<>(keep);

	// initiate all query strings
	this.qCount = "SELECT COUNT(*) FROM \"data_" + this.name + "\"";
	this.qKeys = "SELECT k FROM \"data_" + this.name + "\"";
	this.qValue = "SELECT v, created FROM \"data_" + this.name + "\" WHERE k=?";
	this.qTouch = "UPDATE \"data_" + this.name + "\" SET accessed=? WHERE k=?";
	this.qPut =
	"INSERT INTO \"data_"
	+ this.name
	+ "\" (k, v, created, accessed) VALUES(?,?,?,?) "
	+ "ON CONFLICT (k) DO UPDATE "
	+ "SET v = EXCLUDED.v, created = EXCLUDED.created, accessed = EXCLUDED.accessed";
	this.qInvalidateKey = "DELETE FROM \"data_" + this.name + "\" WHERE k=?";
	this.qInvalidateAll = "DELETE FROM \"data_" + this.name + "\"";
	this.qSum = "SELECT SUM(space) FROM \"data_" + this.name + "\"";
	this.qOrderByAccessed =
	"SELECT k,space,created FROM \"data_" + this.name + "\" ORDER BY accessed";
	this.qStats = "SELECT COUNT(*),SUM(space) FROM \"data_" + this.name + "\"";
	}

	synchronized void open() {
	if (bloomFilter == null) {
	bloomFilter = buildBloomFilter();
	}
	}

	void close() {
	PgSqlHandle h;
	while ((h = handles.poll()) != null) {
	h.close();
	}
	}

	boolean mightContain(K key) {
	BloomFilter<K> b = bloomFilter;
	if (b == null) {
	synchronized (this) {
	b = bloomFilter;
	if (b == null) {
	b = buildBloomFilter();
	bloomFilter = b;
	}
	}
	}
	return b == null \|\| b.mightContain(key);
	}

	private BloomFilter<K> buildBloomFilter() {
	PgSqlHandle c = null;
	try {
	c = acquire();
	try (Statement s = c.conn.createStatement()) {
	if (estimatedSize <= 0) {
	try (ResultSet r = s.executeQuery(qCount)) {
	estimatedSize = r.next() ? r.getInt(1) : 0;
	}
	}

	BloomFilter<K> b = newBloomFilter();
	try (ResultSet r = s.executeQuery(qKeys)) {
	while (r.next()) {
	b.put(keyType.get(r, 1));
	}
	} catch (SQLException e) {
	if (e.getCause() instanceof InvalidClassException) {
	log.warn(
	"Entries cached for {} "
	+ "have an incompatible class and can't be deserialized. "
	+ "Cache is flushed.",
	name);
	invalidateAll();
	} else {
	throw e;
	}
	}
	return b;
	}
	} catch (SQLException e) {
	log.warn("Cannot build BloomFilter for {}: {}", name, e.getMessage());
	c = close(c);
	return null;
	} finally {
	release(c);
	}
	}

	ValueHolder<V> getIfPresent(K key) {
	PgSqlHandle c = null;
	try {
	c = acquire();
	if (c.get == null) {
	c.get = c.conn.prepareStatement(qValue);
	}
	keyType.set(c.get, 1, key);
	try (ResultSet r = c.get.executeQuery()) {
	if (!r.next()) {
	missCount.incrementAndGet();
	return null;
	}

	Timestamp created = r.getTimestamp(2);
	if (expired(created)) {
	invalidate(key);
	missCount.incrementAndGet();
	return null;
	}

	@SuppressWarnings("unchecked")
	V val = (V) PgSqlHandle.deserialize(r.getBytes(1));
	ValueHolder<V> h = new ValueHolder<>(val);
	h.clean = true;
	hitCount.incrementAndGet();
	touch(c, key);
	return h;
	} finally {
	c.get.clearParameters();
	}
	} catch (SQLException e) {
	log.warn("Cannot read cache {} for {}", name, key, e);
	c = close(c);
	return null;
	} finally {
	release(c);
	}
	}

	private boolean expired(Timestamp created) {
	if (expireAfterWrite == 0) {
	return false;
	}
	long age = TimeUtil.nowMs() - created.getTime();
	return 1000 * expireAfterWrite < age;
	}

	private void touch(PgSqlHandle c, K key) throws SQLException {
	if (c.touch == null) {
	c.touch = c.conn.prepareStatement(qTouch);
	}
	try {
	c.touch.setTimestamp(1, TimeUtil.nowTs());
	keyType.set(c.touch, 2, key);
	c.touch.executeUpdate();
	} finally {
	c.touch.clearParameters();
	}
	}

	void put(K key, ValueHolder<V> holder) {
	if (holder.clean) {
	return;
	}

	BloomFilter<K> b = bloomFilter;
	if (b != null) {
	b.put(key);
	bloomFilter = b;
	}

	PgSqlHandle c = null;
	try {
	c = acquire();
	if (c.put == null) {
	c.put = c.conn.prepareStatement(qPut);
	}
	try {
	keyType.set(c.put, 1, key);
	c.put.setObject(2, PgSqlHandle.serialize(holder.value), Types.BINARY);
	c.put.setTimestamp(3, new Timestamp(holder.created));
	c.put.setTimestamp(4, TimeUtil.nowTs());
	c.put.executeUpdate();
	holder.clean = true;
	} finally {
	c.put.clearParameters();
	}
	} catch (SQLException e) {
	log.warn("Cannot put into cache {}", name, e);
	c = close(c);
	} finally {
	release(c);
	}
	}

	void invalidate(K key) {
	PgSqlHandle c = null;
	try {
	c = acquire();
	invalidate(c, key);
	} catch (SQLException e) {
	log.warn("Cannot invalidate cache {}", name, e);
	c = close(c);
	} finally {
	release(c);
	}
	}

	private void invalidate(PgSqlHandle c, K key) throws SQLException {
	if (c.invalidate == null) {
	c.invalidate = c.conn.prepareStatement(qInvalidateKey);
	}
	try {
	keyType.set(c.invalidate, 1, key);
	c.invalidate.executeUpdate();
	} finally {
	c.invalidate.clearParameters();
	}
	}

	void invalidateAll() {
	PgSqlHandle c = null;
	try {
	c = acquire();
	try (Statement s = c.conn.createStatement()) {
	s.executeUpdate(qInvalidateAll);
	}
	bloomFilter = newBloomFilter();
	} catch (SQLException e) {
	log.warn("Cannot invalidate cache {}", name, e);
	c = close(c);
	} finally {
	release(c);
	}
	}

	void prune(Cache<K, ?> mem) {
	PgSqlHandle c = null;
	try {
	c = acquire();
	try (Statement s = c.conn.createStatement()) {
	long used = 0;
	try (ResultSet r = s.executeQuery(qSum)) {
	used = r.next() ? r.getLong(1) : 0;
	}
	if (used <= maxSize) {
	return;
	}

	try (ResultSet r = s.executeQuery(qOrderByAccessed)) {
	while (maxSize < used && r.next()) {
	K key = keyType.get(r, 1);
	Timestamp created = r.getTimestamp(3);
	if (mem.getIfPresent(key) != null && !expired(created)) {
	touch(c, key);
	} else {
	invalidate(c, key);
	used -= r.getLong(2);
	}
	}
	}
	}
	} catch (SQLException e) {
	log.warn("Cannot prune cache {}", name, e);
	c = close(c);
	} finally {
	release(c);
	}
	}

	DiskStats diskStats() {
	long size = 0;
	long space = 0;
	PgSqlHandle c = null;
	try {
	c = acquire();
	try (Statement s = c.conn.createStatement();
	ResultSet r = s.executeQuery(qStats)) {
	if (r.next()) {
	size = r.getLong(1);
	space = r.getLong(2);
	}
	}
	} catch (SQLException e) {
	log.warn("Cannot get DiskStats for {}", name, e);
	c = close(c);
	} finally {
	release(c);
	}
	return new DiskStats(size, space, hitCount.get(), missCount.get());
	}

	private PgSqlHandle acquire() throws SQLException {
	PgSqlHandle h = handles.poll();
	return h != null ? h : new PgSqlHandle(source, name, keyType);
	}

	private void release(PgSqlHandle h) {
	if (h != null && !handles.offer(h)) {
	h.close();
	}
	}

	private PgSqlHandle close(PgSqlHandle h) {
	if (h != null) {
	h.close();
	}
	return null;
	}

	private BloomFilter<K> newBloomFilter() {
	int cnt = Math.max(64 * 1024, 2 * estimatedSize);
	return BloomFilter.create(keyType.funnel(), cnt);
	}
	}