Close the collection iterators (#270)

Author: suresh-prakash

entity-service-impl/src/main/java/org/hypertrace/entity/data/service/EntityDataServiceImpl.java

@@ -11,17 +11,18 @@
 import com.google.protobuf.Message;
 import com.google.protobuf.ServiceException;
 import io.grpc.Channel;
+import io.grpc.Status;
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import org.hypertrace.core.documentstore.CloseableIterator;
 import org.hypertrace.core.documentstore.Collection;
 import org.hypertrace.core.documentstore.Datastore;
 import org.hypertrace.core.documentstore.Document;
@@ -188,14 +189,18 @@ public void getAndUpsertEntities(Entities request, StreamObserver<Entity> respon
         documentMap.put(key, doc);
       }
 
-      List<Entity> existingEntities =
-          Streams.stream(entitiesCollection.bulkUpsertAndReturnOlderDocuments(documentMap))
-              .flatMap(
-                  document -> PARSER.<Entity>parseOrLog(document, Entity.newBuilder()).stream())
-              .map(Entity::toBuilder)
-              .map(builder -> builder.setTenantId(tenantId))
-              .map(Entity.Builder::build)
-              .collect(Collectors.toList());
+      List<Entity> existingEntities;
+      try (final CloseableIterator<Document> iterator =
+          entitiesCollection.bulkUpsertAndReturnOlderDocuments(documentMap)) {
+        existingEntities =
+            Streams.stream(iterator)
+                .flatMap(
+                    document -> PARSER.<Entity>parseOrLog(document, Entity.newBuilder()).stream())
+                .map(Entity::toBuilder)
+                .map(builder -> builder.setTenantId(tenantId))
+                .map(Builder::build)
+                .collect(Collectors.toList());
+      }
 
       existingEntities.forEach(responseObserver::onNext);
       responseObserver.onCompleted();
@@ -299,21 +304,30 @@ public void delete(ByIdRequest request, StreamObserver<Empty> responseObserver)
     Key key =
         this.entityNormalizer.getEntityDocKey(
             tenantId, request.getEntityType(), request.getEntityId());
-    Optional<Entity> existingEntity =
-        this.entityFetcher.getEntitiesByEntityIds(tenantId, List.of(key.toString())).stream()
-            .findFirst();
 
-    if (entitiesCollection.delete(key)) {
-      responseObserver.onNext(Empty.newBuilder().build());
-      responseObserver.onCompleted();
-      existingEntity.ifPresent(
-          entity -> {
-            this.entityCounterMetricSender.sendEntitiesDeleteMetrics(
-                requestContext, request.getEntityType(), List.of(entity));
-            entityChangeEventGenerator.sendDeleteNotification(requestContext, List.of(entity));
-          });
-    } else {
-      responseObserver.onError(new RuntimeException("Could not delete the entity."));
+    try {
+      Optional<Entity> existingEntity =
+          this.entityFetcher.getEntitiesByEntityIds(tenantId, List.of(key.toString())).stream()
+              .findFirst();
+
+      if (entitiesCollection.delete(key)) {
+        responseObserver.onNext(Empty.newBuilder().build());
+        responseObserver.onCompleted();
+        existingEntity.ifPresent(
+            entity -> {
+              this.entityCounterMetricSender.sendEntitiesDeleteMetrics(
+                  requestContext, request.getEntityType(), List.of(entity));
+              entityChangeEventGenerator.sendDeleteNotification(requestContext, List.of(entity));
+            });
+      } else {
+        responseObserver.onError(new RuntimeException("Could not delete the entity"));
+      }
+    } catch (final Exception e) {
+      LOG.error("Unknown error occurred", e);
+      responseObserver.onError(
+          Status.INTERNAL
+              .withDescription("Unknown error occurred")
+              .asRuntimeException(requestContext.buildTrailers()));
     }
   }
 
@@ -326,15 +340,24 @@ public void delete(ByIdRequest request, StreamObserver<Empty> responseObserver)
   @Override
   public void query(Query request, StreamObserver<Entity> responseObserver) {
     logQuery(request);
-    Optional<String> tenantId = RequestContext.CURRENT.get().getTenantId();
+    final RequestContext requestContext = RequestContext.CURRENT.get();
+    Optional<String> tenantId = requestContext.getTenantId();
     if (tenantId.isEmpty()) {
       responseObserver.onError(new ServiceException("Tenant id is missing in the request."));
       return;
     }
 
-    this.entityFetcher
-        .query(DocStoreConverter.transform(tenantId.get(), request, Collections.emptyList()))
-        .forEach(responseObserver::onNext);
+    try {
+      this.entityFetcher
+          .query(DocStoreConverter.transform(tenantId.get(), request, Collections.emptyList()))
+          .forEach(responseObserver::onNext);
+    } catch (final Exception e) {
+      LOG.error("Unknown error occurred", e);
+      responseObserver.onError(
+          Status.INTERNAL
+              .withDescription("Unknown error occurred")
+              .asRuntimeException(requestContext.buildTrailers()));
+    }
 
     responseObserver.onCompleted();
   }
@@ -545,6 +568,7 @@ public void getEnrichedEntityByTypeAndIdentifyingProps(
     String entityId =
         this.entityIdGenerator.generateEntityId(
             tenantId, request.getEntityType(), request.getIdentifyingAttributesMap());
+
     searchByIdAndStreamSingleResponse(
         tenantId,
         entityId,
@@ -565,32 +589,32 @@ public void mergeAndUpsertEntity(
       return;
     }
 
-    Entity receivedEntity = this.entityNormalizer.normalize(tenantId, request.getEntity());
-    Optional<Entity> existingEntity =
-        getExistingEntity(tenantId, receivedEntity.getEntityType(), receivedEntity.getEntityId());
+    try {
+      Entity receivedEntity = this.entityNormalizer.normalize(tenantId, request.getEntity());
+      Optional<Entity> existingEntity =
+          getExistingEntity(tenantId, receivedEntity.getEntityType(), receivedEntity.getEntityId());
 
-    boolean rejectUpsertForConditionMismatch =
-        existingEntity
-            .map(
-                entity ->
-                    !this.upsertConditionMatcher.matches(entity, request.getUpsertCondition()))
-            .orElse(false);
-
-    if (rejectUpsertForConditionMismatch) {
-      // There's an existing entity and the update doesn't meet the condition, return existing
-      responseObserver.onNext(
-          MergeAndUpsertEntityResponse.newBuilder().setEntity(existingEntity.get()).build());
-      responseObserver.onCompleted();
-    } else {
-      // There's either a new entity or a valid update to upsert
-      Entity entityToUpsert =
+      boolean rejectUpsertForConditionMismatch =
           existingEntity
-              .map(Entity::toBuilder)
-              .map(Entity.Builder::clearCreatedTime)
-              .map(builder -> builder.mergeFrom(receivedEntity))
-              .map(Builder::build)
-              .orElse(receivedEntity);
-      try {
+              .map(
+                  entity ->
+                      !this.upsertConditionMatcher.matches(entity, request.getUpsertCondition()))
+              .orElse(false);
+
+      if (rejectUpsertForConditionMismatch) {
+        // There's an existing entity and the update doesn't meet the condition, return existing
+        responseObserver.onNext(
+            MergeAndUpsertEntityResponse.newBuilder().setEntity(existingEntity.get()).build());
+        responseObserver.onCompleted();
+      } else {
+        // There's either a new entity or a valid update to upsert
+        Entity entityToUpsert =
+            existingEntity
+                .map(Entity::toBuilder)
+                .map(Entity.Builder::clearCreatedTime)
+                .map(builder -> builder.mergeFrom(receivedEntity))
+                .map(Builder::build)
+                .orElse(receivedEntity);
         Entity upsertedEntity = this.upsertEntity(tenantId, entityToUpsert);
         responseObserver.onNext(
             MergeAndUpsertEntityResponse.newBuilder().setEntity(upsertedEntity).build());
@@ -604,9 +628,11 @@ public void mergeAndUpsertEntity(
             List.of(upsertedEntity));
         entityChangeEventGenerator.sendChangeNotification(
             requestContext, existingEntities, List.of(upsertedEntity));
-      } catch (IOException e) {
-        responseObserver.onError(e);
       }
+    } catch (final Exception e) {
+      LOG.error("Unknown error occurred", e);
+      responseObserver.onError(
+          Status.INTERNAL.withDescription("Unknown error occurred").asRuntimeException());
     }
   }
 
@@ -716,23 +742,32 @@ private <T extends Message> void searchByIdAndStreamSingleResponse(
     String docId = this.entityNormalizer.getEntityDocKey(tenantId, entityType, entityId).toString();
     query.setFilter(new Filter(Filter.Op.EQ, EntityServiceConstants.ID, docId));
 
-    Iterator<Document> result = collection.search(query);
     List<T> entities = new ArrayList<>();
-    while (result.hasNext()) {
-      PARSER
-          .<T>parseOrLog(result.next(), builder.clone())
-          .map(
-              entity -> {
-                // Populate the tenant id field with the tenant id that's received for backward
-                // compatibility.
-                Descriptors.FieldDescriptor fieldDescriptor =
-                    entity.getDescriptorForType().findFieldByName("tenant_id");
-                if (fieldDescriptor != null) {
-                  return (T) entity.toBuilder().setField(fieldDescriptor, tenantId).build();
-                }
-                return entity;
-              })
-          .ifPresent(entities::add);
+    try (final CloseableIterator<Document> result = collection.search(query)) {
+      while (result.hasNext()) {
+        PARSER
+            .<T>parseOrLog(result.next(), builder.clone())
+            .map(
+                entity -> {
+                  // Populate the tenant id field with the tenant id that's received for backward
+                  // compatibility.
+                  Descriptors.FieldDescriptor fieldDescriptor =
+                      entity.getDescriptorForType().findFieldByName("tenant_id");
+                  if (fieldDescriptor != null) {
+                    return (T) entity.toBuilder().setField(fieldDescriptor, tenantId).build();
+                  }
+                  return entity;
+                })
+            .ifPresent(entities::add);
+      }
+    } catch (final IOException e) {
+      final String message =
+          String.format(
+              "Unable to search for tenant: %s, entityType: %s, entityId: %s",
+              tenantId, entityType, entityId);
+      LOG.warn(message, e);
+      responseObserver.onError(Status.INTERNAL.withDescription(message).asRuntimeException());
+      return;
     }
 
     if (LOG.isDebugEnabled()) {
@@ -754,7 +789,8 @@ private <T extends Message> void searchByIdAndStreamSingleResponse(
     }
   }
 
-  private List<Entity> getExistingEntities(String tenantId, java.util.Collection<Entity> entities) {
+  private List<Entity> getExistingEntities(String tenantId, java.util.Collection<Entity> entities)
+      throws IOException {
     List<String> docIds = entities.stream().map(this::getDocId).collect(Collectors.toList());
     return this.entityFetcher.getEntitiesByDocIds(tenantId, docIds);
   }
@@ -765,7 +801,8 @@ private String getDocId(Entity entity) {
         .toString();
   }
 
-  private Optional<Entity> getExistingEntity(String tenantId, String entityType, String entityId) {
+  private Optional<Entity> getExistingEntity(String tenantId, String entityType, String entityId)
+      throws IOException {
     String docId = this.entityNormalizer.getEntityDocKey(tenantId, entityType, entityId).toString();
     return this.entityFetcher.getEntitiesByDocIds(tenantId, List.of(docId)).stream().findFirst();
   }
@@ -778,23 +815,29 @@ private void searchByQueryAndStreamRelationships(
       org.hypertrace.core.documentstore.Query query,
       StreamObserver<EntityRelationship> responseObserver,
       String tenantId) {
-    List<EntityRelationship> relationships =
-        Streams.stream(relationshipsCollection.search(query))
-            .flatMap(
-                document ->
-                    PARSER
-                        .<EntityRelationship>parseOrLog(document, EntityRelationship.newBuilder())
-                        .stream())
-            .map(EntityRelationship::toBuilder)
-            .map(builder -> builder.setTenantId(tenantId))
-            .map(EntityRelationship.Builder::build)
-            .peek(responseObserver::onNext)
-            .collect(Collectors.toList());
+    try (final CloseableIterator<Document> documentIterator =
+        relationshipsCollection.search(query)) {
+      List<EntityRelationship> relationships =
+          Streams.stream(documentIterator)
+              .flatMap(
+                  document ->
+                      PARSER
+                          .<EntityRelationship>parseOrLog(document, EntityRelationship.newBuilder())
+                          .stream())
+              .map(EntityRelationship::toBuilder)
+              .map(builder -> builder.setTenantId(tenantId))
+              .map(EntityRelationship.Builder::build)
+              .peek(responseObserver::onNext)
+              .collect(Collectors.toList());
 
-    if (LOG.isDebugEnabled()) {
       LOG.debug("Docstore query has returned the result: {}", relationships);
+      responseObserver.onCompleted();
+    } catch (final IOException e) {
+      final String message =
+          String.format("Unknown error occurred for query: %s on tenant: %s", query, tenantId);
+      LOG.warn(message, e);
+      responseObserver.onError(Status.INTERNAL.withDescription(message).asRuntimeException());
     }
-    responseObserver.onCompleted();
   }
 
   private void logQuery(Object query) {

entity-service-impl/src/main/java/org/hypertrace/entity/data/service/IdentifyingAttributeCache.java

@@ -6,12 +6,15 @@
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
 import com.google.common.collect.Streams;
+import io.grpc.Status;
+import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
+import org.hypertrace.core.documentstore.CloseableIterator;
 import org.hypertrace.core.documentstore.Collection;
 import org.hypertrace.core.documentstore.Datastore;
 import org.hypertrace.core.documentstore.Document;
@@ -48,11 +51,17 @@ private Map<String, List<AttributeType>> loadEntityTypes(String tenantId) {
         new Filter(
             Op.IN, EntityServiceConstants.TENANT_ID, TenantUtils.getTenantHierarchy(tenantId)));
 
-    return Streams.stream(entityTypesCollection.search(query))
-        .flatMap(this::buildEntityType)
-        .collect(
-            Collectors.toUnmodifiableMap(
-                EntityType::getName, this::getIdentifyingAttributesFromType));
+    try (final CloseableIterator<Document> iterator = entityTypesCollection.search(query)) {
+      return Streams.stream(iterator)
+          .flatMap(this::buildEntityType)
+          .collect(
+              Collectors.toUnmodifiableMap(
+                  EntityType::getName, this::getIdentifyingAttributesFromType));
+    } catch (final IOException e) {
+      throw Status.INTERNAL
+          .withDescription("Unable to fetch entity types for tenant: " + tenantId)
+          .asRuntimeException();
+    }
   }
 
   private Stream<EntityType> buildEntityType(Document doc) {