From 27dd9a4effdbf9761b569cf2dae8a924291a438b Mon Sep 17 00:00:00 2001
From: Greg Burd <gburd@onshape.com>
Date: Wed, 24 Jan 2018 09:28:19 -0500
Subject: [PATCH] Include the post commit/abort function execution time that
 happens within/associated-with the scope of a UOW.  Fix reset method to check
 both maps for updates.

---
 .../net/helenus/core/AbstractEntityDraft.java    | 11 ++++++++++-
 .../net/helenus/core/AbstractUnitOfWork.java     | 16 +++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/main/java/net/helenus/core/AbstractEntityDraft.java b/src/main/java/net/helenus/core/AbstractEntityDraft.java
index e1dd4f0..259f103 100644
--- a/src/main/java/net/helenus/core/AbstractEntityDraft.java
+++ b/src/main/java/net/helenus/core/AbstractEntityDraft.java
@@ -151,10 +151,19 @@ public abstract class AbstractEntityDraft<E> implements Drafted<E> {
     return this.<T>reset(this.<T>methodNameFor(getter), desiredValue);
   }
 
+  private <T> T fetch(String key) {
+      T value = (T) backingMap.get(key);
+
+      if (value == null) {
+          value = (T) entityMap.get(key);
+      }
+      return value;
+  }
+
   public <T> boolean reset(String key, T desiredValue) {
     if (key != null && desiredValue != null) {
       @SuppressWarnings("unchecked")
-      T currentValue = (T) backingMap.get(key);
+      T currentValue = (T) this.<T>fetch(key);
       if (currentValue == null || !currentValue.equals(desiredValue)) {
         set(key, desiredValue);
         return true;
diff --git a/src/main/java/net/helenus/core/AbstractUnitOfWork.java b/src/main/java/net/helenus/core/AbstractUnitOfWork.java
index 52134d2..9a1f7af 100644
--- a/src/main/java/net/helenus/core/AbstractUnitOfWork.java
+++ b/src/main/java/net/helenus/core/AbstractUnitOfWork.java
@@ -200,9 +200,6 @@ public abstract class AbstractUnitOfWork<E extends Exception>
         f.apply();
       }
     }
-    if (LOG.isInfoEnabled()) {
-      LOG.info(logTimers(what));
-    }
   }
 
   @Override
@@ -365,7 +362,6 @@ public abstract class AbstractUnitOfWork<E extends Exception>
     }
 
     if (!canCommit) {
-      elapsedTime.stop();
 
       if (parent == null) {
 
@@ -376,11 +372,16 @@ public abstract class AbstractUnitOfWork<E extends Exception>
                 uow -> {
                   applyPostCommitFunctions("aborted", abortThunks);
                 });
+
+        elapsedTime.stop();
+        if (LOG.isInfoEnabled()) {
+          LOG.info(logTimers("aborted"));
+        }
+
       }
 
       return new PostCommitFunction(this, null, null, false);
     } else {
-      elapsedTime.stop();
       committed = true;
       aborted = false;
 
@@ -404,6 +405,11 @@ public abstract class AbstractUnitOfWork<E extends Exception>
                     new HelenusException(
                         "Futures must be resolved before their unit of work has committed/aborted.")));
 
+        elapsedTime.stop();
+        if (LOG.isInfoEnabled()) {
+          LOG.info(logTimers("committed"));
+        }
+
         return new PostCommitFunction(this, null, null, true);
       } else {