Merge pull request #507 from hubmapconsortium/yuanzhou/cache-improvement

Cache for complete entity and normalzied entity

Author: yuanzhou

src/app.py

@@ -163,15 +163,15 @@ def http_internal_server_error(e):
 if MEMCACHED_MODE:
     try:
         # Use client pool to maintain a pool of already-connected clients for improved performance
-        # The uwsgi config launches the app across multiple threads (2) inside each process (4), making essentially 8 processes
+        # The uwsgi config launches the app across multiple threads (8) inside each process (32), making essentially 256 processes
         # Set the connect_timeout and timeout to avoid blocking the process when memcached is slow, defaults to "forever"
         # connect_timeout: seconds to wait for a connection to the memcached server
         # timeout: seconds to wait for send or reveive calls on the socket connected to memcached
         # Use the ignore_exc flag to treat memcache/network errors as cache misses on calls to the get* methods
         # Set the no_delay flag to sent TCP_NODELAY (disable Nagle's algorithm to improve TCP/IP networks and decrease the number of packets)
         # If you intend to use anything but str as a value, it is a good idea to use a serializer
         memcached_client_instance = PooledClient(app.config['MEMCACHED_SERVER'], 
-                                                 max_pool_size = 8,
+                                                 max_pool_size = 256,
                                                  connect_timeout = 1,
                                                  timeout = 30,
                                                  ignore_exc = True, 
@@ -369,7 +369,7 @@ def flush_all_cache():
 Parameters
 ----------
 id : str
-    The HuBMAP ID (e.g. HBM123.ABCD.456) or UUID of target entity (Donor/Dataset/Sample/Upload)
+    The HuBMAP ID (e.g. HBM123.ABCD.456) or UUID of target entity (Donor/Dataset/Sample/Upload/Collection/Publication)
 
 Returns
 -------
@@ -381,12 +381,8 @@ def flush_cache(id):
     msg = ''
 
     if MEMCACHED_MODE:
-        msg = f'No cache found from Memcached for entity {id}'
-        cache_key = f'{MEMCACHED_PREFIX}{id}'
-
-        if memcached_client_instance.get(cache_key) is not None:
-            memcached_client_instance.delete(cache_key)
-            msg = f'The cached data has been deleted from Memcached for entity {id}'
+        delete_cache(id)
+        msg = f'The cached data has been deleted from Memcached for entity {id}'
     else:
         msg = 'No caching is being used because Memcached mode is not enabled at all'
 
@@ -633,14 +629,16 @@ def get_entity_by_id(id):
     # Otherwise query against uuid-api and neo4j to get the entity dict if the id exists
     entity_dict = query_target_entity(id, token)
     normalized_entity_type = entity_dict['entity_type']
-    # To verify if a Collection is public, it is necessary to have its Datasets, which
-    # are populated as triggered data.  So pull back the complete entity for
-    # _get_entity_visibility() to check.
+
+    # Get the generated complete entity result from cache if exists
+    # Otherwise re-generate on the fly
     complete_dict = schema_manager.get_complete_entity_result(token, entity_dict)
 
     # Determine if the entity is publicly visible base on its data, only.
-    entity_scope = _get_entity_visibility(  normalized_entity_type=normalized_entity_type
-                                            ,entity_dict=complete_dict)
+    # To verify if a Collection is public, it is necessary to have its Datasets, which
+    # are populated as triggered data.  So pull back the complete entity for
+    # _get_entity_visibility() to check.
+    entity_scope = _get_entity_visibility(normalized_entity_type=normalized_entity_type, entity_dict=complete_dict)
 
     # Initialize the user as authorized if the data is public.  Otherwise, the
     # user is not authorized and credentials must be checked.
@@ -1441,34 +1439,18 @@ def update_entity(id):
         if (return_all_properties is not None) and (return_all_properties.lower() == 'true'):
             properties_to_skip = []
 
-    # Generate the filtered or complete entity dict to send back
-    complete_dict = schema_manager.get_complete_entity_result(user_token, merged_updated_dict, properties_to_skip)
-
-    # Will also filter the result based on schema
-    normalized_complete_dict = schema_manager.normalize_entity_result_for_response(complete_dict)
-
     # Remove the cached entities if Memcached is being used
+    # DO NOT update the cache with new entity dict because the returned dict from PUT (some properties maybe skipped)
+    # can be different from the one generated by GET call
     if MEMCACHED_MODE:
-        # Delete the old cache data of this entity
-        # DO NOT update the cache with new entity dict because the returned dict from PUT (some properties maybe skipped)
-        # can be different from the one generated by GET call
-        cache_key = f'{MEMCACHED_PREFIX}{id}'
-        memcached_client_instance.delete(cache_key)
-
-        logger.info(f"Deleted cache of key: {cache_key} after entity update via PUT call")
-
-        # Also delete the cache of all the direct descendants (children)
-        # Otherwise they'll have old cached data for the `direct_ancestor` (Sample) `direct_ancestors` (Dataset) fields
-        # Note: must use uuid in the Neo4j query
-        children_uuid_list = schema_neo4j_queries.get_children(neo4j_driver_instance, entity_dict['uuid'] , 'uuid')
-        
-        logger.info(f"Also delete the cache of all the direct descendants (children) of {id} if exist")
+        delete_cache(id)
 
-        for child_uuid in children_uuid_list:
-            cache_key = f'{MEMCACHED_PREFIX}{child_uuid}'
-            memcached_client_instance.delete(cache_key)
+    # Do not return the updated dict to avoid computing overhead - 7/14/2023 by Zhou
+    # # Generate the complete entity dict
+    # complete_dict = schema_manager.get_complete_entity_result(user_token, merged_updated_dict, properties_to_skip)
 
-            logger.info(f"Deleted direct descendant cache of key: {cache_key}")
+    # # Will also filter the result based on schema
+    # normalized_complete_dict = schema_manager.normalize_entity_result_for_response(complete_dict)
 
     # Also reindex the updated entity node in elasticsearch via search-api
     if entity_dict['entity_type'] in ['Collection']:
@@ -1480,9 +1462,13 @@ def update_entity(id):
         logger.log(logging.INFO
                    ,f"Re-indexing for creation of {entity_dict['entity_type']}"
                     f" with UUID {entity_dict['uuid']}")
+
         reindex_entity(entity_dict['uuid'], user_token)
 
-    return jsonify(normalized_complete_dict)
+    # Do not return the updated dict to avoid computing overhead - 7/14/2023 by Zhou
+    # return jsonify(normalized_complete_dict)
+
+    return jsonify({'message': f"{normalized_entity_type} of {id} has been updated"})
 
 
 """
@@ -3372,11 +3358,9 @@ def sankey_data():
         if memcached_client_instance.get(cache_key) is not None:
             dataset_sankey_list = memcached_client_instance.get(cache_key)
 
-    current_datetime = datetime.now()
-
     if not dataset_sankey_list:
         if MEMCACHED_MODE:
-            logger.info(f'Sankey data cache not found or expired. Making a new data fetch at time {current_datetime}')
+            logger.info(f'Sankey data cache not found or expired. Making a new data fetch at time {datetime.now()}')
 
         # Call to app_neo4j_queries to prepare and execute the database query
         sankey_info = app_neo4j_queries.get_sankey_info(neo4j_driver_instance)
@@ -3415,7 +3399,7 @@ def sankey_data():
             # Cache the result
             memcached_client_instance.set(cache_key, dataset_sankey_list, expire = SchemaConstants.MEMCACHED_TTL)
     else:
-        logger.info(f'Using the cached sankey data at time {current_datetime}')
+        logger.info(f'Using the cached sankey data at time {datetime.now()}')
 
     return jsonify(dataset_sankey_list)
 
@@ -4409,7 +4393,7 @@ def after_update(normalized_entity_type, user_token, entity_dict):
 
 
 """
-Get target entity dict for the given id
+Get target entity dict from Neo4j query for the given id
 
 Parameters
 ----------
@@ -4425,20 +4409,20 @@ def after_update(normalized_entity_type, user_token, entity_dict):
 """
 def query_target_entity(id, user_token):
     entity_dict = None
-
-    cache_key = f'{MEMCACHED_PREFIX}{id}'
+    cache_result = None
 
-    if MEMCACHED_MODE:
+    if MEMCACHED_MODE and MEMCACHED_PREFIX and memcached_client_instance:
+        # If this id is hubmap_id rather than uuid, there won't be a cache
+        # Only uuid is used in the cache key
+        cache_key = f'{MEMCACHED_PREFIX}_neo4j_{id}'
         # Memcached returns None if no cached data or expired
-        entity_dict = memcached_client_instance.get(cache_key)
+        cache_result = memcached_client_instance.get(cache_key)
 
-    current_datetime = datetime.now()
-
-    # Use the cached data if found and still valid
-    # Otherwise, make a fresh query and add to cache
-    if entity_dict is None:
-        if MEMCACHED_MODE:
-            logger.info(f'Cache not found or expired. Making a new query to retrieve {id} at time {current_datetime}')
+    # Use the cached data if the id is an uuid and we found a valid cache
+    # Otherwise, either the id is a hubmap_id or we don't have a cache for it even if it's uuid
+    if cache_result is None:
+        if MEMCACHED_MODE and MEMCACHED_PREFIX and memcached_client_instance:
+            logger.info(f'Neo4j entity cache of {id} not found or expired at time {datetime.now()}')
 
         try:
             """
@@ -4462,15 +4446,31 @@ def query_target_entity(id, user_token):
 
             # Get the target uuid if all good
             uuid = hubmap_ids['hm_uuid']
-            entity_dict = schema_neo4j_queries.get_entity(neo4j_driver_instance, uuid)
 
-            # The uuid exists via uuid-api doesn't mean it's also in Neo4j
-            if not entity_dict:
-                not_found_error(f"Entity of id: {id} not found in Neo4j")
-            
-            if MEMCACHED_MODE:
-                # Cache the result
-                memcached_client_instance.set(cache_key, entity_dict, expire = SchemaConstants.MEMCACHED_TTL)
+            # Look up the cache again by the uuid since we only use uuid in the cache key
+            if MEMCACHED_MODE and MEMCACHED_PREFIX and memcached_client_instance:
+                cache_key = f'{MEMCACHED_PREFIX}_neo4j_{uuid}'
+                cache_result = memcached_client_instance.get(cache_key)
+
+                if cache_result is None:
+                    logger.info(f'Neo4j entity cache of {uuid} not found or expired at time {datetime.now()}')
+
+                    # Make a new query against neo4j
+                    entity_dict = schema_neo4j_queries.get_entity(neo4j_driver_instance, uuid)
+
+                    # The uuid exists via uuid-api doesn't mean it also exists in Neo4j
+                    if not entity_dict:
+                        not_found_error(f"Entity of id: {uuid} not found in Neo4j")
+                    
+                    logger.info(f'Creating neo4j entity result cache of {uuid} at time {datetime.now()}')
+
+                    cache_key = f'{MEMCACHED_PREFIX}_neo4j_{uuid}'
+                    memcached_client_instance.set(cache_key, entity_dict, expire = SchemaConstants.MEMCACHED_TTL)
+                else:
+                    logger.info(f'Using neo4j entity cache of {uuid} at time {datetime.now()}')
+                    logger.debug(entity_dict)
+
+                    entity_dict = cache_result
         except requests.exceptions.RequestException as e:
             # Due to the use of response.raise_for_status() in schema_manager.get_hubmap_ids()
             # we can access the status codes from the exception
@@ -4483,10 +4483,12 @@ def query_target_entity(id, user_token):
             else:
                 internal_server_error(e.response.text)
     else:
-        logger.info(f'Using the cache data of entity {id} at time {current_datetime}')
+        logger.info(f'Using neo4j entity cache of {id} at time {datetime.now()}')
         logger.debug(entity_dict)
 
-    # Final return
+        entity_dict = cache_result
+
+    # One final return
     return entity_dict
 
 
@@ -4501,6 +4503,42 @@ def require_json(request):
         bad_request_error("A json body and appropriate Content-Type header are required")
 
 
+
+"""
+Delete the cached data of all possible keys used for the given entity id
+
+Parameters
+----------
+id : str
+    The HuBMAP ID (e.g. HBM123.ABCD.456) or UUID of target entity (Donor/Dataset/Sample/Upload/Collection/Publication)
+"""
+def delete_cache(id):
+    if MEMCACHED_MODE:
+        # First delete the target entity cache
+        entity_dict = query_target_entity(id, get_internal_token())
+        entity_uuid = entity_dict['uuid']
+
+        # If the target entity is Sample (`direct_ancestor`) or Dataset/Publication (`direct_ancestors`)
+        # Delete the cache of all the direct descendants (children)
+        child_uuids = schema_neo4j_queries.get_children(neo4j_driver_instance, entity_uuid , 'uuid')
+
+        # If the target entity is Collection, delete the cache for each of its associated 
+        # Datasets and Publications (via [:IN_COLLECTION] relationship) as well as just Publications (via [:USES_DATA] relationship)
+        collection_dataset_uuids = schema_neo4j_queries.get_collection_associated_datasets(neo4j_driver_instance, entity_uuid , 'uuid')
+
+        # If the target entity is Upload, delete the cache for each of its associated Datasets (via [:IN_UPLOAD] relationship)
+        upload_dataset_uuids = schema_neo4j_queries.get_upload_datasets(neo4j_driver_instance, entity_uuid , 'uuid')
+
+        # If the target entity is Datasets/Publication, delete the associated Collections cache, Upload cache
+        collection_uuids = schema_neo4j_queries.get_dataset_collections(neo4j_driver_instance, entity_uuid , 'uuid')
+        collection_dict = schema_neo4j_queries.get_publication_associated_collection(neo4j_driver_instance, entity_uuid)
+        upload_dict = schema_neo4j_queries.get_dataset_upload(neo4j_driver_instance, entity_uuid)
+
+        # We only use uuid in the cache key acorss all the cache types
+        uuids_list = [entity_uuid] + child_uuids + collection_dataset_uuids + upload_dataset_uuids + collection_uuids + [collection_dict['uuid']] + [upload_dict['uuid']]
+        schema_manager.delete_memcached_cache(uuids_list)
+
+
 """
 Make a call to each search-api instance to reindex this entity node in elasticsearch