sharded-counter-java/src/ShardedCounterV2.java

   1 import com.google.appengine.api.datastore.DatastoreService;
   2 import com.google.appengine.api.datastore.DatastoreServiceFactory;
   3 import com.google.appengine.api.datastore.Entity;
   4 import com.google.appengine.api.datastore.EntityNotFoundException;
   5 import com.google.appengine.api.datastore.Key;
   6 import com.google.appengine.api.datastore.KeyFactory;
   7 import com.google.appengine.api.datastore.Query;
   8 import com.google.appengine.api.datastore.Transaction;
   9 import com.google.appengine.api.memcache.Expiration;
  10 import com.google.appengine.api.memcache.MemcacheService;
  11 import com.google.appengine.api.memcache.MemcacheService.SetPolicy;
  12 import com.google.appengine.api.memcache.MemcacheServiceFactory;
  13
  14 import java.util.Random;
  15
  16 /**
  17  * A counter which can be incremented rapidly.
  18  *
  19  * Capable of incrementing the counter and increasing the number of shards. When
  20  * incrementing, a random shard is selected to prevent a single shard from being
  21  * written to too frequently. If increments are being made too quickly, increase
  22  * the number of shards to divide the load. Performs datastore operations using
  23  * the low level datastore API.
  24  */
  25 public class ShardedCounterV2 {
  26
  27   interface Counter {
  28     static final String KIND = "Counter";
  29     static final String SHARD_COUNT = "shard_count";
  30   }
  31
  32   interface CounterShard {
  33     static final String COUNT = "count";
  34     static final String KIND = "CounterShard";
  35   }
  36
  37   private static final DatastoreService ds = DatastoreServiceFactory
  38       .getDatastoreService();
  39
  40   /**
  41    * Default number of shards.
  42    */
  43   private static final long INITIAL_SHARDS = 5;
  44
  45   /**
  46    * The name of this counter.
  47    */
  48   private String counterName;
  49
  50   /**
  51    * A random number generating, for distributing writes across shards.
  52    */
  53   private Random generator = new Random();
  54
  55   /**
  56    * The counter shard kind for this counter.
  57    */
  58   private String kind;
  59
  60   private MemcacheService mc = MemcacheServiceFactory.getMemcacheService();
  61
  62   public ShardedCounterV2(String counterName) {
  63     this.counterName = counterName;
  64     kind = CounterShard.KIND + "_" + counterName;
  65   }
  66
  67   /**
  68    * Increase the number of shards for a given sharded counter. Will never
  69    * decrease the number of shards.
  70    *
  71    * @param count Number of new shards to build and store
  72    */
  73   public void addShards(int count) {
  74     Key counterKey = KeyFactory.createKey(Counter.KIND, counterName);
  75     incrementPropertyTx(counterKey, Counter.SHARD_COUNT, count, INITIAL_SHARDS
  76         + count);
  77   }
  78
  79   /**
  80    * Retrieve the value of this sharded counter.
  81    *
  82    * @return Summed total of all shards' counts
  83    */
  84   public long getCount() {
  85     Long value = (Long) mc.get(kind);
  86     if (value != null) {
  87       return value;
  88     }
  89
  90     long sum = 0;
  91     Query query = new Query(kind);
  92     for (Entity shard : ds.prepare(query).asIterable()) {
  93       sum += (Long) shard.getProperty(CounterShard.COUNT);
  94     }
  95     mc.put(kind, sum, Expiration.byDeltaSeconds(60),
  96         SetPolicy.ADD_ONLY_IF_NOT_PRESENT);
  97
  98     return sum;
  99   }
 100
 101   /**
 102    * Increment the value of this sharded counter.
 103    */
 104   public void increment() {
 105     // Find how many shards are in this counter.
 106     int numShards = (int) getShardCount();
 107
 108     // Choose the shard randomly from the available shards.
 109     long shardNum = generator.nextInt(numShards);
 110
 111     Key shardKey = KeyFactory.createKey(kind, Long.toString(shardNum));
 112     incrementPropertyTx(shardKey, CounterShard.COUNT, 1, 1);
 113     mc.increment(kind, 1);
 114   }
 115
 116   /**
 117    * Get the number of shards in this counter.
 118    *
 119    * @return shard count
 120    */
 121   private long getShardCount() {
 122     try {
 123       Key counterKey = KeyFactory.createKey(Counter.KIND, counterName);
 124       Entity counter = ds.get(counterKey);
 125       return (Long) counter.getProperty(Counter.SHARD_COUNT);
 126     } catch (EntityNotFoundException ignore) {
 127       return INITIAL_SHARDS;
 128     }
 129   }
 130
 131   /**
 132    * Increment datastore property value inside a transaction. If the entity with
 133    * the provided key does not exist, instead create an entity with the supplied
 134    * initial property value.
 135    *
 136    * @param key the entity key to update or create
 137    * @param prop the property name to be incremented
 138    * @param increment the amount by which to increment
 139    * @param initialValue the value to use if the entity does not exist
 140    */
 141   private void incrementPropertyTx(Key key, String prop, long increment,
 142       long initialValue) {
 143     Transaction tx = ds.beginTransaction();
 144     Entity thing;
 145     long value;
 146     try {
 147       thing = ds.get(tx, key);
 148       value = (Long) thing.getProperty(prop) + increment;
 149     } catch (EntityNotFoundException e) {
 150       thing = new Entity(key);
 151       value = initialValue;
 152     }
 153     thing.setUnindexedProperty(prop, value);
 154     ds.put(tx, thing);
 155     tx.commit();
 156   }
 157 }