Generic JPA sharded counter for Google App Engine

One of the drawback of The Google App Engine datastore is the rate at which it can handle update for a single entity or an entity group. The datastore documentation indicates that the maximum rates is arround 5 write operations per second for the same entity or entity group. To overpass this limitation, Google recommend to use horizontal partitioning by using sharded counter. The App Engine documentation provides a simple (an non transactional) JDO implementation of sharded counters. Let see how we can build an reusable and transactional JPA-based sharded counter for the Google App Engine.

The Counter Class

The Entity bellow allow us to implement sharded counters for almost anything :

count :

this attribute will handle the counter value for this sharded counter

refId :

the ID of the Entity this sharded counter is used for

entityClass :

the Entity class this sharded counter is used for

type :

what this counter is counting

@Entity
@SuppressWarnings("serial")
public class Counter implements Serializable {
    
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;

    private Long count;
    private Long refId;
    private int type;
    private String entityClass;
    
    public Long getId() {
        return id;
    }
    public void setId(Long id) {
        this.id = id;
    }
    public Long getCount() {
        return count;
    }
    public void setCount(Long count) {
        this.count = count;
    }
    public Long getRefId() {
        return refId;
    }
    public void setRefId(Long refId) {
        this.refId = refId;
    }
    public String getEntityClass() {
        return entityClass;
    }
    public void setEntityClass(String entityClass) {
        this.entityClass = entityClass;
    }
    public int getType() {
        return type;
    }
    public void setType(int type) {
        this.type = type;
    }

}

The service and how to increment a counter

In order to increment the number of time a web page is viewed, we would like to do something like :

incrementCounter(WebPage.class, webPageId, WebPage.VIEWS)

Where webPageId is the id of the WebPage Entity and WebPage.VIEW a constant.

The incrementCounter method will work as follow :

1. Defines a MAX_TRIES values to store the maximum number of time we will try to update an existing sharded counter
2. Retrives the list of the sharded counter already persisted for the given type
3. If none exists, a new sharded counter with a value of 1 is persisted for the given type, the method returns
4. Else, one sharded counter is picked up at random and its value is incremented
5. If the update fails, the number of remaining tries is decremented
6. If there is no try left, a new sharded counter with a value of 1 is persisted for the given type, the method returns
7. Else, start again at step 2

/**
 * Sum all the counter for the given type and entity
 * @param c
 * @param refId
 * @param type
 * @return
 */
protected long counterSum(Class<?> c, Long refId, int type) {
    long sum = 0;
    List<Counter> counters = getCounters(c, refId, type);
    for (Counter counter : counters) {
        sum += counter.getCount();
    }
    return sum;
}

/**
 * Get all the counter for the given type and entity
 * @param c
 * @param refId
 * @param type
 * @return
 */
private List<Counter> getCounters(Class<?> c, Long refId, int type) {
    Map<String, Object> params = new HashMap<String, Object>();
    params.put("refId", refId);
    params.put("type", type);
    params.put("entityClass", c.getSimpleName());
    return list(Counter.class, 
        "SELECT c FROM Counter c WHERE refId = :refId AND type = :type AND entityClass = :entityClass", 
        params);
}

protected void incrementCounter(Class<?> c, Long refId, int type) {
    modifyCounter(c, refId, type, 1);
}

protected void decrementCounter(Class<?> c, Long refId, int type) {
    modifyCounter(c, refId, type, -1);
}

/**
 * Modify the counter value for the given type and entity
 * @param c
 * @param refId
 * @param type
 * @param step
 */
protected void modifyCounter(Class<?> c, Long refId, int type, int step) {
    int tries = MAX_TRIES;
    EntityManager em = getEntityManager();
    while (true) {
        try {
            List<Counter> counters = getCounters(c, refId, type);
            if (counters.size() == 0) {
                newCounter(c, refId, type, step);
                break;
            }
            try {
                em.getTransaction().begin();
                Random generator = new Random();
                int counterNum = generator.nextInt(counters.size());
                Counter counter = counters.get(counterNum);
                counter.setCount(counter.getCount() + step);
                em.merge(counter);
                em.getTransaction().commit();
                break;
            } finally {
                if (em != null) {
                    if (em.getTransaction().isActive()) {
                        em.getTransaction().rollback();
                    }
                }
            }
        } catch (ConcurrentModificationException cme) {
            if (--tries == 0) {
                newCounter(c, refId, type, step);
                break;
            }
        }
    }
}

private void newCounter(Class<?> c, Long refId, int type, int step) {
    EntityManager em = null;
    try {
        em = getEntityManager();
        Counter counter = new Counter();
        counter.setCount(Long.valueOf(step));
        counter.setEntityClass(c.getSimpleName());
        counter.setRefId(refId);
        counter.setType(type);
        em.getTransaction().begin();
        em.persist(counter);
        em.getTransaction().commit();
    } finally {
        if (em != null) {
            if (em.getTransaction().isActive()) {
                em.getTransaction().rollback();
            }
        }
    }
}


protected final <R> List<R> list(Class<R> c, String query, Map<String, Object> parameters) {
    EntityManager em = getEntityManager();
    Query select = em.createQuery(query);
    for (String key : parameters.keySet()) {
        select.setParameter(key, parameters.get(key));
    }
    List<R> list = select.getResultList();
    return list;
}

Enjoy!