Skip to content

Commit

Permalink
Refactor out External*Resolver classes, rename ResourceReaders.Extern…
Browse files Browse the repository at this point in the history
…alDelegate and ModuleKeys.External to MessageTransport
  • Loading branch information
HT154 committed Oct 16, 2024
1 parent b3c3641 commit 98ec98f
Show file tree
Hide file tree
Showing 7 changed files with 312 additions and 227 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
/**
* Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
*
* Licensed under the Apache 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
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.pkl.core.externalReader;

import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Random;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import org.pkl.core.SecurityManager;
import org.pkl.core.SecurityManagerException;
import org.pkl.core.messaging.MessageTransport;
import org.pkl.core.messaging.MessageTransports;
import org.pkl.core.messaging.Messages.ListModulesRequest;
import org.pkl.core.messaging.Messages.ListModulesResponse;
import org.pkl.core.messaging.Messages.ReadModuleRequest;
import org.pkl.core.messaging.Messages.ReadModuleResponse;
import org.pkl.core.messaging.ProtocolException;
import org.pkl.core.module.PathElement;

public class ExternalModuleResolver {
private final MessageTransport transport;
private final long evaluatorId;
private final Map<URI, Future<String>> readResponses = new ConcurrentHashMap<>();
private final Map<URI, Future<List<PathElement>>> listResponses = new ConcurrentHashMap<>();

public ExternalModuleResolver(MessageTransport transport, long evaluatorId) {
this.transport = transport;
this.evaluatorId = evaluatorId;
}

public List<PathElement> listElements(SecurityManager securityManager, URI uri)
throws IOException, SecurityManagerException {
securityManager.checkResolveModule(uri);
return doListElements(uri);
}

public boolean hasElement(SecurityManager securityManager, URI uri)
throws SecurityManagerException {
securityManager.checkResolveModule(uri);
try {
doReadModule(uri);
return true;
} catch (IOException e) {
return false;
}
}

public String resolveModule(SecurityManager securityManager, URI uri)
throws IOException, SecurityManagerException {
securityManager.checkResolveModule(uri);
return doReadModule(uri);
}

private String doReadModule(URI moduleUri) throws IOException {
return MessageTransports.resolveFuture(
readResponses.computeIfAbsent(
moduleUri,
(uri) -> {
var future = new CompletableFuture<String>();
var request = new ReadModuleRequest(new Random().nextLong(), evaluatorId, uri);
try {
transport.send(
request,
(response) -> {
if (response instanceof ReadModuleResponse resp) {
if (resp.getError() != null) {
future.completeExceptionally(new IOException(resp.getError()));
} else if (resp.getContents() != null) {
future.complete(resp.getContents());
} else {
future.complete("");
}
} else {
future.completeExceptionally(new ProtocolException("unexpected response"));
}
});
} catch (ProtocolException | IOException e) {
future.completeExceptionally(e);
}
return future;
}));
}

private List<PathElement> doListElements(URI baseUri) throws IOException {
return MessageTransports.resolveFuture(
listResponses.computeIfAbsent(
baseUri,
(uri) -> {
var future = new CompletableFuture<List<PathElement>>();
var request = new ListModulesRequest(new Random().nextLong(), evaluatorId, uri);
try {
transport.send(
request,
(response) -> {
if (response instanceof ListModulesResponse resp) {
if (resp.getError() != null) {
future.completeExceptionally(new IOException(resp.getError()));
} else {
future.complete(
Objects.requireNonNullElseGet(resp.getPathElements(), List::of));
}
} else {
future.completeExceptionally(new ProtocolException("unexpected response"));
}
});
} catch (ProtocolException | IOException e) {
future.completeExceptionally(e);
}
return future;
}));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
/**
* Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
*
* Licensed under the Apache 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
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.pkl.core.externalReader;

import java.io.IOException;
import java.net.URI;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Random;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import org.pkl.core.SecurityManager;
import org.pkl.core.SecurityManagerException;
import org.pkl.core.messaging.MessageTransport;
import org.pkl.core.messaging.MessageTransports;
import org.pkl.core.messaging.Messages.Bytes;
import org.pkl.core.messaging.Messages.ListResourcesRequest;
import org.pkl.core.messaging.Messages.ListResourcesResponse;
import org.pkl.core.messaging.Messages.ReadResourceRequest;
import org.pkl.core.messaging.Messages.ReadResourceResponse;
import org.pkl.core.messaging.ProtocolException;
import org.pkl.core.module.PathElement;
import org.pkl.core.resource.Resource;

public class ExternalResourceResolver {
private final MessageTransport transport;
private final long evaluatorId;
private final Map<URI, Future<Bytes>> readResponses = new ConcurrentHashMap<>();
private final Map<URI, Future<List<PathElement>>> listResponses = new ConcurrentHashMap<>();

public ExternalResourceResolver(MessageTransport transport, long evaluatorId) {
this.transport = transport;
this.evaluatorId = evaluatorId;
}

public Optional<Object> read(URI uri) throws IOException {
var result = doRead(uri);
return Optional.of(new Resource(uri, result.getBytes()));
}

public boolean hasElement(org.pkl.core.SecurityManager securityManager, URI elementUri)
throws SecurityManagerException {
securityManager.checkResolveResource(elementUri);
try {
doRead(elementUri);
return true;
} catch (IOException e) {
return false;
}
}

public List<PathElement> listElements(SecurityManager securityManager, URI baseUri)
throws IOException, SecurityManagerException {
securityManager.checkResolveResource(baseUri);
return doListElements(baseUri);
}

public List<PathElement> doListElements(URI baseUri) throws IOException {
return MessageTransports.resolveFuture(
listResponses.computeIfAbsent(
baseUri,
(uri) -> {
var future = new CompletableFuture<List<PathElement>>();
var request = new ListResourcesRequest(new Random().nextLong(), evaluatorId, uri);
try {
transport.send(
request,
(response) -> {
if (response instanceof ListResourcesResponse resp) {
if (resp.getError() != null) {
future.completeExceptionally(new IOException(resp.getError()));
} else {
future.complete(
Objects.requireNonNullElseGet(resp.getPathElements(), List::of));
}
} else {
future.completeExceptionally(new ProtocolException("unexpected response"));
}
});
} catch (ProtocolException | IOException e) {
future.completeExceptionally(e);
}
return future;
}));
}

public Bytes doRead(URI baseUri) throws IOException {
return MessageTransports.resolveFuture(
readResponses.computeIfAbsent(
baseUri,
(uri) -> {
var future = new CompletableFuture<Bytes>();
var request = new ReadResourceRequest(new Random().nextLong(), evaluatorId, uri);
try {
transport.send(
request,
(response) -> {
if (response instanceof ReadResourceResponse resp) {
if (resp.getError() != null) {
future.completeExceptionally(new IOException(resp.getError()));
} else if (resp.getContents() != null) {
future.complete(resp.getContents());
} else {
future.complete(new Bytes(new byte[0]));
}
} else {
future.completeExceptionally(new ProtocolException("unexpected response"));
}
});
} catch (ProtocolException | IOException e) {
future.completeExceptionally(e);
}
return future;
}));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import java.util.List;
import java.util.Optional;
import java.util.ServiceLoader;
import org.pkl.core.externalReader.ExternalModuleResolver;
import org.pkl.core.externalReader.ExternalReaderProcess;
import org.pkl.core.externalReader.ExternalReaderProcessException;
import org.pkl.core.util.ErrorMessages;
Expand Down Expand Up @@ -255,20 +256,20 @@ private static final class External implements ModuleKeyFactory {
private final String scheme;
private final ExternalReaderProcess process;
private final long evaluatorId;
private ModuleKeys.External.Resolver resolver;
private ExternalModuleResolver resolver;

public External(String scheme, ExternalReaderProcess process, long evaluatorId) {
this.scheme = scheme;
this.process = process;
this.evaluatorId = evaluatorId;
}

private ModuleKeys.External.Resolver getResolver() throws ExternalReaderProcessException {
private ExternalModuleResolver getResolver() throws ExternalReaderProcessException {
if (resolver != null) {
return resolver;
}

resolver = new ModuleKeys.External.Resolver(process.getTransport(), evaluatorId);
resolver = new ExternalModuleResolver(process.getTransport(), evaluatorId);
return resolver;
}

Expand All @@ -282,7 +283,7 @@ public Optional<ModuleKey> create(URI uri)
ErrorMessages.create("externalReaderDoesNotSupportScheme", "module", scheme));
}

return Optional.of(ModuleKeys.external(uri, spec, getResolver()));
return Optional.of(ModuleKeys.messageTransport(uri, spec, getResolver()));
}

@Override
Expand Down
Loading

0 comments on commit 98ec98f

Please sign in to comment.