killbill-memoizeit

diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java
index 72287d4..f06589f 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntityDaoBase.java
@@ -47,13 +47,15 @@ public abstract class EntityDaoBase<M extends EntityModelDao<E>, E extends Entit
 
     @Override
     public void create(final M entity, final InternalCallContext context) throws U {
-        transactionalSqlDao.execute(getCreateEntitySqlDaoTransactionWrapper(entity, context));
+        final M refreshedEntity = transactionalSqlDao.execute(getCreateEntitySqlDaoTransactionWrapper(entity, context));
+        // Populate the caches only after the transaction has been committed, in case of rollbacks
+        transactionalSqlDao.populateCaches(refreshedEntity);
     }
 
-    protected EntitySqlDaoTransactionWrapper<Void> getCreateEntitySqlDaoTransactionWrapper(final M entity, final InternalCallContext context) {
-        return new EntitySqlDaoTransactionWrapper<Void>() {
+    protected EntitySqlDaoTransactionWrapper<M> getCreateEntitySqlDaoTransactionWrapper(final M entity, final InternalCallContext context) {
+        return new EntitySqlDaoTransactionWrapper<M>() {
             @Override
-            public Void inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
+            public M inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final EntitySqlDao<M, E> transactional = entitySqlDaoWrapperFactory.become(realSqlDao);
 
                 if (checkEntityAlreadyExists(transactional, entity, context)) {
@@ -62,7 +64,7 @@ public abstract class EntityDaoBase<M extends EntityModelDao<E>, E extends Entit
                 final M refreshedEntity = createAndRefresh(transactional, entity, context);
 
                 postBusEventFromTransaction(entity, refreshedEntity, ChangeType.INSERT, entitySqlDaoWrapperFactory, context);
-                return null;
+                return refreshedEntity;
             }
         };
     }

diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoTransactionalJdbiWrapper.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoTransactionalJdbiWrapper.java
index e70fcba..b16601d 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoTransactionalJdbiWrapper.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoTransactionalJdbiWrapper.java
@@ -49,6 +49,10 @@ public class EntitySqlDaoTransactionalJdbiWrapper {
         this.internalCallContextFactory = internalCallContextFactory;
     }
 
+    public <M extends EntityModelDao> void populateCaches(final M refreshedEntity) {
+        EntitySqlDaoWrapperInvocationHandler.populateCaches(cacheControllerDispatcher, refreshedEntity);
+    }
+
     class JdbiTransaction<ReturnType, M extends EntityModelDao<E>, E extends Entity> implements Transaction<ReturnType, EntitySqlDao<M, E>> {
 
         private final Handle h;

diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
index efc6b37..c747bfa 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/EntitySqlDaoWrapperInvocationHandler.java
@@ -349,8 +349,11 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
         }
     }
 
-    private void populateCacheOnGetByIdInvocation(M model) {
+    private void populateCacheOnGetByIdInvocation(final M model) {
+        populateCaches(cacheControllerDispatcher, model);
+    }
 
+    public static void populateCaches(final CacheControllerDispatcher cacheControllerDispatcher, final EntityModelDao model) {
         final CacheController<String, Long> cacheRecordId = cacheControllerDispatcher.getCacheController(CacheType.RECORD_ID);
         cacheRecordId.putIfAbsent(getKey(model.getId().toString(), CacheType.RECORD_ID, model.getTableName()), model.getRecordId());
 
@@ -368,7 +371,7 @@ public class EntitySqlDaoWrapperInvocationHandler<S extends EntitySqlDao<M, E>, 
         }
     }
 
-    private String getKey(final String rawKey, final CacheType cacheType, final TableName tableName) {
+    private static String getKey(final String rawKey, final CacheType cacheType, final TableName tableName) {
         return cacheType.isKeyPrefixedWithTableName() ?
                tableName + CacheControllerDispatcher.CACHE_KEY_SEPARATOR + rawKey :
                rawKey;