harden client references so the server is not crashed due to client bugs

Author: Myron Scott

tc-client/src/main/java/com/tc/object/ClientEntityManagerImpl.java

@@ -57,7 +57,6 @@
 import com.tc.text.MapListPrettyPrint;
 import com.tc.text.PrettyPrintable;
 import com.tc.util.Assert;
-import com.tc.util.Util;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 
@@ -409,6 +408,23 @@ private <M extends EntityMessage, R extends EntityResponse> EntityClientEndpoint
     Assert.assertNotNull("Can't lookup null entity descriptor", instance);
     final EntityDescriptor fetchDescriptor = EntityDescriptor.createDescriptorForFetch(entity, version, instance);
     EntityClientEndpointImpl<M, R> resolvedEndpoint = null;
+    // make sure this close hook is only ever called once
+    CloseHookCallable closeCall = new CloseHookCallable() {
+      boolean released = false;
+      @Override
+      public EntityException call() {
+        if (!released) {
+          released = true;
+          try {
+            internalRelease(entity, fetchDescriptor);
+          } catch (EntityException e) {
+            // We aren't expecting there to be any problems releasing an entity in the close hook but return the exception if there is an issue.
+            return e;
+          }
+        }
+        return null;
+      }
+    };
     try {
       byte[] raw = internalRetrieve(fetchDescriptor);
       ByteBuffer br = ByteBuffer.wrap(raw);
@@ -418,21 +434,8 @@ private <M extends EntityMessage, R extends EntityResponse> EntityClientEndpoint
       br.get(config);
       // We can only fail to get the config if we threw an exception.
       Assert.assertTrue(null != raw);
-      // We managed to retrieve the config so create the end-point.
-      // Note that we will need to call release on this descriptor, when it is closed, prior to running the closeHook we
-      //  were given so combine these ideas to pass in one runnable.
-      Runnable compoundRunnable = new Runnable() {
-        @Override
-        public void run() {
-          try {
-            internalRelease(entity, fetchDescriptor);
-          } catch (EntityException e) {
-            // We aren't expecting there to be any problems releasing an entity in the close hook so we will just log and re-throw.
-            Util.printLogAndRethrowError(e, logger);
-          }
-        }
-      };
-      resolvedEndpoint = new EntityClientEndpointImpl<>(entity, version, EntityDescriptor.createDescriptorForInvoke(fetch, instance), this, config, codec, compoundRunnable, this.endpointCloser);
+
+      resolvedEndpoint = new EntityClientEndpointImpl<>(entity, version, EntityDescriptor.createDescriptorForInvoke(fetch, instance), this, config, codec, closeCall, this.endpointCloser);
 
       if (this.objectStoreMap.putIfAbsent(instance, resolvedEndpoint) != null) {
         throw Assert.failure("Attempt to add an object that already exists: Object of class " + resolvedEndpoint.getClass()
@@ -442,14 +445,14 @@ public void run() {
       throw notfound;
     } catch (EntityException e) {
       // Release the entity and re-throw to the higher level.
-      internalRelease(entity, fetchDescriptor);
-      // NOTE:  Since we are throwing, we are not responsible for calling the given closeHook.
+      e.addSuppressed(closeCall.call());
+
       throw e;
     } catch (Throwable t) {
       // This is the unexpected case so clean up and re-throw as a RuntimeException
       // Clean up any client-side or server-side state regarding this failed connection.
-      internalRelease(entity, fetchDescriptor);
-      // NOTE:  Since we are throwing, we are not responsible for calling the given closeHook.
+      t.addSuppressed(closeCall.call());
+
       throw Throwables.propagate(t);
     }
 

tc-client/src/main/java/com/tc/object/CloseHookCallable.java

@@ -0,0 +1,30 @@
+/*
+ *
+ *  The contents of this file are subject to the Terracotta Public 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://terracotta.org/legal/terracotta-public-license.
+ *
+ *  Software distributed under the License is distributed on an "AS IS" basis,
+ *  WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License for
+ *  the specific language governing rights and limitations under the License.
+ *
+ *  The Covered Software is Terracotta Core.
+ *
+ *  The Initial Developer of the Covered Software is
+ *  Terracotta, Inc., a Software AG company
+ *
+ */
+package com.tc.object;
+
+import java.util.concurrent.Callable;
+import org.terracotta.exception.EntityException;
+
+/**
+ *
+ */
+public interface CloseHookCallable extends Callable<EntityException> {
+  @Override
+  EntityException call();
+}

tc-client/src/main/java/com/tc/object/EntityClientEndpointImpl.java

@@ -40,6 +40,7 @@
 import org.slf4j.LoggerFactory;
 
 import static com.tc.object.SafeInvocationCallback.safe;
+import org.terracotta.exception.EntityException;
 
 
 public class EntityClientEndpointImpl<M extends EntityMessage, R extends EntityResponse> implements EntityClientEndpoint<M, R> {
@@ -52,7 +53,7 @@ public class EntityClientEndpointImpl<M extends EntityMessage, R extends EntityR
   private final EntityID entityID;
   private final long version;
   private final MessageCodec<M, R> codec;
-  private final Runnable closeHook;
+  private final CloseHookCallable closeHook;
   private final ExecutorService closer;
   private EndpointDelegate<R> delegate;
   private boolean isOpen;
@@ -66,7 +67,7 @@ public class EntityClientEndpointImpl<M extends EntityMessage, R extends EntityR
    * @param entityConfiguration Opaque byte[] describing how to configure the entity to be built on top of this end-point.
    * @param closeHook A Runnable which will be run last when the end-point is closed.
    */
-  public EntityClientEndpointImpl(EntityID eid, long version, EntityDescriptor instance, InvocationHandler invocationHandler, byte[] entityConfiguration, MessageCodec<M, R> codec, Runnable closeHook, ExecutorService closer) {
+  public EntityClientEndpointImpl(EntityID eid, long version, EntityDescriptor instance, InvocationHandler invocationHandler, byte[] entityConfiguration, MessageCodec<M, R> codec, CloseHookCallable closeHook, ExecutorService closer) {
     this.entityID = eid;
     this.version = version;
     this.invokeDescriptor = instance;
@@ -96,14 +97,14 @@ EntityDescriptor getEntityDescriptor() {
   @Override
   public byte[] getEntityConfiguration() {
     // This is harmless while closed but shouldn't be called so check open.
-    checkEndpointOpen();
+    checkEndpointOpen(false);
     return configuration;
   }
 
   @Override
   public void setDelegate(EndpointDelegate<R> delegate) {
     // This is harmless while closed but shouldn't be called so check open.
-    checkEndpointOpen();
+    checkEndpointOpen(false);
     Assert.assertNull(this.delegate);
     this.delegate = delegate;
   }
@@ -124,7 +125,7 @@ public InFlightStats getStatistics() {
   @Override
   public Invocation<R> message(M message) {
     // We can't create new invocations when the endpoint is closed.
-    checkEndpointOpen();
+    checkEndpointOpen(false);
     return new InvocationImpl(message);
   }
 
@@ -176,12 +177,17 @@ public byte[] getExtendedReconnectData() {
   @Override
   public void close() {
     // We can't close twice.
-    checkEndpointOpen();
+    try {
+      checkEndpointOpen(true);
+    } catch (IllegalStateException alreadyClosed) {
+      LOGGER.info("Trying to close an endpoint which has already been closed", alreadyClosed);
+      return;
+    }
     if (this.closeHook != null) {
-      this.closeHook.run();
+      EntityException e = this.closeHook.call();
+      LOGGER.info("Exception occured during close", e);
+      // log and swallow this exception closing
     }
-    // We also need to invalidate ourselves so we don't continue allowing new messages through when disconnecting.
-    this.isOpen = false;
   }
 
   @Override
@@ -220,9 +226,12 @@ public void didCloseUnexpectedly() {
     }
   }
 
-  private void checkEndpointOpen() {
+  private synchronized void checkEndpointOpen(boolean close) {
     if (!this.isOpen) {
       throw new IllegalStateException("Endpoint closed");
     }
+    if (close) {
+      this.isOpen = false;
+    }
   }
 }

tc-server/src/main/java/com/tc/objectserver/entity/ManagedEntityImpl.java

@@ -976,16 +976,28 @@ private void getEntity(ServerEntityRequest getEntityRequest, ResultCapture respo
     if (this.isDestroyed) {
       response.failure(ServerException.createNotFoundException(getID()));
     } else {
-      if (canDelete) {
-        clientReferenceCount += 1;
-        Assert.assertTrue(clientReferenceCount > 0);
-      }
       // The FETCH can only come directly from a client so we can down-cast.
       ClientID clientID = getEntityRequest.getNodeID();
       ClientDescriptorImpl descriptor = new ClientDescriptorImpl(clientID, getEntityRequest.getClientInstance());
       boolean added = clientEntityStateManager.addReference(descriptor, this.fetchID);
+      
+      if (canDelete) {
+        if (added) {
+          clientReferenceCount += 1;
+          Assert.assertTrue(clientReferenceCount > 0);
+        } else {
+          // this should never happen.  Ther server is expecting sane things from the client. 
+          logger.info("the client has attempted to fetch the same entity instance twice {}", descriptor);
+        }
+      }
+      
       if (this.isInActiveState) {
-        Assert.assertTrue(added);
+        if (!added) {
+          //  this should never happen if the client is sane but if it does, exception the client
+          //  don't crash the server
+          response.failure(ServerException.createReferencedException(id));
+          return;
+        }
         // Fire the event that the client fetched the entity.
         this.eventCollector.clientDidFetchEntity(clientID, this.id, this.consumerID, getEntityRequest.getClientInstance());
         // finally notify the entity that it was fetched
@@ -1006,10 +1018,18 @@ private void getEntity(ServerEntityRequest getEntityRequest, ResultCapture respo
               handler.handleReconnect(descriptor, extendedData);
             }
           } catch (ReconnectRejectedException rejected) {
+            Assert.assertTrue(clientEntityStateManager.removeReference(descriptor));
+            if (canDelete) {
+              clientReferenceCount -= 1;
+            }
             response.failure(ServerException.createReconnectRejected(getID(), rejected));
             return;
           } catch (Exception e) {
 //  something happened during reconnection, force a disconnection, see ProcessTransactionHandler.disconnectClientDueToFailure for handling
+            Assert.assertTrue(clientEntityStateManager.removeReference(descriptor));
+            if (canDelete) {
+              clientReferenceCount -= 1;
+            }
             logger.warn("unexpected exception.  rejecting reconnection of " + descriptor.getNodeID() + " to " + this.id, e);
             response.failure(ServerException.createReconnectRejected(getID(), new ReconnectRejectedException(e.getMessage(), e)));
             return;
@@ -1031,17 +1051,28 @@ private void releaseEntity(ServerEntityRequest request, ResultCapture response)
     if (this.isDestroyed) {
       response.failure(ServerException.createNotFoundException(this.getID()));
     } else {
-      if (canDelete) {
-        clientReferenceCount -= 1;
-        Assert.assertTrue(clientReferenceCount >= 0);
-      }
-
       ClientID clientID = request.getNodeID();
       ClientDescriptorImpl clientInstance = new ClientDescriptorImpl(clientID, request.getClientInstance());
       boolean removed = clientEntityStateManager.removeReference(clientInstance);
 
+      if (canDelete) {
+        if (removed) {
+          clientReferenceCount -= 1;
+          Assert.assertTrue(clientReferenceCount >= 0);
+        } else {
+        // this should never happen.  the expectation is that the client will send sane things to the server.  Will assert if active
+          logger.info("client has tried to release and entity that is not fetched {}", clientInstance);
+        }
+      }
+
       if (this.isInActiveState) {
-        Assert.assertTrue(removed);
+        if (!removed) {
+          //  this should never happen if the client is sane but if it does, exception the client
+          //  don't crash the server
+          response.failure(ServerException.createNotFoundException(id));
+          return;
+        }
+        
         this.activeServerEntity.disconnected(clientInstance);
         // Fire the event that the client released the entity.
         this.eventCollector.clientDidReleaseEntity(clientID, this.id, this.consumerID, request.getClientInstance());
@@ -1077,6 +1108,8 @@ public Runnable promoteEntity() throws ConfigurationException {
 // up on passives
     if (canDelete) {
       this.clientReferenceCount = 0;
+      // the entity references should have been cleared
+      Assert.assertTrue(clientEntityStateManager.verifyNoEntityReferences(fetchID));
     } else {
       Assert.assertEquals(this.clientReferenceCount, ManagedEntityImpl.UNDELETABLE_ENTITY);
     }

tc-server/src/test/java/com/tc/objectserver/entity/ManagedEntityImplTest.java

@@ -173,6 +173,7 @@ public void setUp() throws Exception {
     messageSelf = mock(Sink.class);
     when(clientEntityStateManager.addReference(any(ClientDescriptorImpl.class), any(FetchID.class))).thenReturn(Boolean.TRUE);
     when(clientEntityStateManager.removeReference(any(ClientDescriptorImpl.class))).thenReturn(Boolean.TRUE);
+    when(clientEntityStateManager.verifyNoEntityReferences(any())).thenReturn(Boolean.TRUE);
     eventCollector = new ManagementTopologyEventCollector(mock(IMonitoringProducer.class));
     // We will start this in a passive state, as the general test case.
     boolean isInActiveState = false;