Skip to content

Commit

Permalink
cache public rules
Browse files Browse the repository at this point in the history
  • Loading branch information
@artsiomkorzun
artsiomkorzun committed Mar 14, 2024
1 parent c650edd commit f2a4e17
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 8 deletions.
26 changes: 24 additions & 2 deletions src/main/java/com/epam/aidial/core/service/PublicationService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import com.epam.aidial.core.data.MetadataBase;
import com.epam.aidial.core.data.Publication;
import com.epam.aidial.core.data.ResourceFolderMetadata;
import com.epam.aidial.core.data.ResourceItemMetadata;
import com.epam.aidial.core.data.ResourceType;
import com.epam.aidial.core.data.ResourceUrl;
import com.epam.aidial.core.data.Rule;
Expand All @@ -17,6 +18,7 @@
import com.fasterxml.jackson.core.type.TypeReference;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.commons.lang3.tuple.Pair;

import java.util.Collection;
import java.util.HashMap;
Expand All @@ -25,6 +27,7 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.LongSupplier;
import java.util.function.Supplier;
import javax.annotation.Nullable;
Expand Down Expand Up @@ -53,6 +56,7 @@ public class PublicationService {

private static final Set<ResourceType> ALLOWED_RESOURCES = Set.of(ResourceType.FILE, ResourceType.CONVERSATION, ResourceType.PROMPT);

private final AtomicReference<Pair<Long, Map<String, List<Rule>>>> cachedRules = new AtomicReference<>();
private final EncryptionService encryption;
private final ResourceService resources;
private final BlobStorage files;
Expand All @@ -78,7 +82,7 @@ public boolean hasPublicAccess(ProxyContext context, ResourceDescription resourc
return false;
}

Map<String, List<Rule>> rules = decodeRules(resources.getResource(PUBLIC_RULES));
Map<String, List<Rule>> rules = getCachedRules();
Map<String, Boolean> cache = new HashMap<>();

return evaluate(context, resource, rules, cache);
Expand All @@ -89,7 +93,7 @@ public void filterForbidden(ProxyContext context, ResourceDescription folder, Re
return;
}

Map<String, List<Rule>> rules = decodeRules(resources.getResource(PUBLIC_RULES));
Map<String, List<Rule>> rules = getCachedRules();
Map<String, Boolean> cache = new HashMap<>();
cache.put(ruleUrl(folder), true);

Expand Down Expand Up @@ -473,6 +477,24 @@ private String encodeReviewBucket(ResourceDescription bucket, String id) {
return encryption.encrypt(path);
}

private Map<String, List<Rule>> getCachedRules() {
ResourceItemMetadata meta = resources.getResourceMetadata(PUBLIC_RULES);
long key = (meta == null) ? Long.MIN_VALUE : meta.getUpdatedAt();
Pair<Long, Map<String, List<Rule>>> current = cachedRules.get();

if (current == null || current.getKey() != key) {
Pair<ResourceItemMetadata, String> resource = resources.getResourceWithMetadata(PUBLIC_RULES);
Pair<Long, Map<String, List<Rule>>> next = (resource == null)
? Pair.of(Long.MIN_VALUE, decodeRules(null))
: Pair.of(resource.getKey().getUpdatedAt(), decodeRules(resource.getValue()));

cachedRules.compareAndSet(current, next);
current = next;
}

return current.getValue();
}

private static boolean evaluate(ProxyContext context,
ResourceDescription resource,
Map<String, List<Rule>> rules,
Expand Down
36 changes: 30 additions & 6 deletions src/main/java/com/epam/aidial/core/service/ResourceService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import lombok.Getter;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Pair;
import org.jclouds.blobstore.domain.Blob;
import org.jclouds.blobstore.domain.BlobMetadata;
import org.jclouds.blobstore.domain.PageSet;
Expand Down Expand Up @@ -114,7 +115,7 @@ public void close() {
public MetadataBase getMetadata(ResourceDescription descriptor, String token, int limit, boolean recursive) {
return descriptor.isFolder()
? getFolderMetadata(descriptor, token, limit, recursive)
: getItemMetadata(descriptor);
: getResourceMetadata(descriptor);
}

private ResourceFolderMetadata getFolderMetadata(ResourceDescription descriptor, String token, int limit, boolean recursive) {
Expand Down Expand Up @@ -156,7 +157,11 @@ private ResourceFolderMetadata getFolderMetadata(ResourceDescription descriptor,
return new ResourceFolderMetadata(descriptor, resources).setNextToken(set.getNextMarker());
}

private ResourceItemMetadata getItemMetadata(ResourceDescription descriptor) {
public ResourceItemMetadata getResourceMetadata(ResourceDescription descriptor) {
if (descriptor.isFolder()) {
throw new IllegalArgumentException("Resource folder: " + descriptor.getUrl());
}

String redisKey = redisKey(descriptor);
String blobKey = blobKey(descriptor);
Result result = redisGet(redisKey, false);
Expand Down Expand Up @@ -187,12 +192,12 @@ public boolean hasResource(ResourceDescription descriptor) {
}

@Nullable
public String getResource(ResourceDescription descriptor) {
return getResource(descriptor, true);
public Pair<ResourceItemMetadata, String> getResourceWithMetadata(ResourceDescription descriptor) {
return getResourceWithMetadata(descriptor, true);
}

@Nullable
public String getResource(ResourceDescription descriptor, boolean lock) {
public Pair<ResourceItemMetadata, String> getResourceWithMetadata(ResourceDescription descriptor, boolean lock) {
String redisKey = redisKey(descriptor);
Result result = redisGet(redisKey, true);

Expand All @@ -208,7 +213,26 @@ public String getResource(ResourceDescription descriptor, boolean lock) {
}
}

return result.exists ? result.body : null;
if (result.exists) {
ResourceItemMetadata metadata = new ResourceItemMetadata(descriptor)
.setCreatedAt(result.createdAt)
.setUpdatedAt(result.updatedAt);

return Pair.of(metadata, result.body);
}

return null;
}

@Nullable
public String getResource(ResourceDescription descriptor) {
return getResource(descriptor, true);
}

@Nullable
public String getResource(ResourceDescription descriptor, boolean lock) {
Pair<ResourceItemMetadata, String> result = getResourceWithMetadata(descriptor, lock);
return (result == null) ? null : result.getRight();
}

public ResourceItemMetadata putResource(ResourceDescription descriptor, String body, boolean overwrite) {
Expand Down

0 comments on commit f2a4e17

Please sign in to comment.