| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you 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 |
| * |
| * http://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.apache.beam.runners.fnexecution.artifact; |
| |
| import java.io.FileInputStream; |
| import java.io.IOException; |
| import java.io.InputStream; |
| import java.nio.channels.Channels; |
| import java.nio.charset.StandardCharsets; |
| import java.nio.file.Files; |
| import java.nio.file.Path; |
| import java.nio.file.Paths; |
| import java.util.ArrayList; |
| import java.util.Collections; |
| import java.util.Iterator; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.Set; |
| import java.util.concurrent.CompletableFuture; |
| import java.util.concurrent.ExecutionException; |
| import java.util.concurrent.ExecutorService; |
| import java.util.concurrent.Executors; |
| import java.util.concurrent.TimeUnit; |
| import java.util.stream.Collectors; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ArtifactChunk; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ArtifactMetadata; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.CommitManifestRequest; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.GetArtifactRequest; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.GetManifestRequest; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.GetManifestResponse; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.Manifest; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ProxyManifest; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.ProxyManifest.Location; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.PutArtifactMetadata; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.PutArtifactRequest; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactApi.PutArtifactResponse; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactRetrievalServiceGrpc; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactRetrievalServiceGrpc.ArtifactRetrievalServiceBlockingStub; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactRetrievalServiceGrpc.ArtifactRetrievalServiceStub; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc.ArtifactStagingServiceBlockingStub; |
| import org.apache.beam.model.jobmanagement.v1.ArtifactStagingServiceGrpc.ArtifactStagingServiceStub; |
| import org.apache.beam.runners.fnexecution.GrpcFnServer; |
| import org.apache.beam.runners.fnexecution.InProcessServerFactory; |
| import org.apache.beam.sdk.io.FileSystems; |
| import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString; |
| import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.ManagedChannel; |
| import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.inprocess.InProcessChannelBuilder; |
| import org.apache.beam.vendor.grpc.v1p21p0.io.grpc.stub.StreamObserver; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Strings; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.hash.Hashing; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.io.ByteStreams; |
| import org.junit.After; |
| import org.junit.Assert; |
| import org.junit.Before; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.TemporaryFolder; |
| import org.junit.runner.RunWith; |
| import org.junit.runners.JUnit4; |
| |
| /** |
| * Tests for {@link BeamFileSystemArtifactStagingService} and {@link |
| * BeamFileSystemArtifactRetrievalService}. |
| */ |
| @RunWith(JUnit4.class) |
| public class BeamFileSystemArtifactServicesTest { |
| private static final int DATA_1KB = 1 << 10; |
| private GrpcFnServer<BeamFileSystemArtifactStagingService> stagingServer; |
| private BeamFileSystemArtifactStagingService stagingService; |
| private GrpcFnServer<BeamFileSystemArtifactRetrievalService> retrievalServer; |
| private BeamFileSystemArtifactRetrievalService retrievalService; |
| private ArtifactStagingServiceStub stagingStub; |
| private ArtifactStagingServiceBlockingStub stagingBlockingStub; |
| private ArtifactRetrievalServiceStub retrievalStub; |
| private ArtifactRetrievalServiceBlockingStub retrievalBlockingStub; |
| private Path stagingDir; |
| private Path originalDir; |
| @Rule public TemporaryFolder tempFolder = new TemporaryFolder(); |
| |
| @Before |
| public void setUp() throws Exception { |
| stagingService = new BeamFileSystemArtifactStagingService(); |
| stagingServer = |
| GrpcFnServer.allocatePortAndCreateFor(stagingService, InProcessServerFactory.create()); |
| ManagedChannel stagingChannel = |
| InProcessChannelBuilder.forName(stagingServer.getApiServiceDescriptor().getUrl()).build(); |
| stagingStub = ArtifactStagingServiceGrpc.newStub(stagingChannel); |
| stagingBlockingStub = ArtifactStagingServiceGrpc.newBlockingStub(stagingChannel); |
| |
| retrievalService = new BeamFileSystemArtifactRetrievalService(); |
| retrievalServer = |
| GrpcFnServer.allocatePortAndCreateFor(retrievalService, InProcessServerFactory.create()); |
| ManagedChannel retrievalChannel = |
| InProcessChannelBuilder.forName(retrievalServer.getApiServiceDescriptor().getUrl()).build(); |
| retrievalStub = ArtifactRetrievalServiceGrpc.newStub(retrievalChannel); |
| retrievalBlockingStub = ArtifactRetrievalServiceGrpc.newBlockingStub(retrievalChannel); |
| |
| originalDir = tempFolder.newFolder("original").toPath(); |
| stagingDir = tempFolder.newFolder("staging").toPath(); |
| } |
| |
| @After |
| public void tearDown() throws Exception { |
| if (stagingServer != null) { |
| stagingServer.close(); |
| } |
| if (stagingService != null) { |
| stagingService.close(); |
| } |
| if (retrievalServer != null) { |
| retrievalServer.close(); |
| } |
| if (retrievalService != null) { |
| retrievalService.close(); |
| } |
| } |
| |
| private void putArtifact(String stagingSessionToken, final String filePath, final String fileName) |
| throws Exception { |
| CompletableFuture<Boolean> complete = new CompletableFuture<>(); |
| StreamObserver<PutArtifactRequest> outputStreamObserver = |
| stagingStub.putArtifact( |
| new StreamObserver<PutArtifactResponse>() { |
| @Override |
| public void onNext(PutArtifactResponse putArtifactResponse) { |
| // Do nothing. |
| } |
| |
| @Override |
| public void onError(Throwable throwable) { |
| throwable.printStackTrace(); |
| Assert.fail("OnError should never be called."); |
| } |
| |
| @Override |
| public void onCompleted() { |
| complete.complete(Boolean.TRUE); |
| } |
| }); |
| outputStreamObserver.onNext( |
| PutArtifactRequest.newBuilder() |
| .setMetadata( |
| PutArtifactMetadata.newBuilder() |
| .setMetadata(ArtifactMetadata.newBuilder().setName(fileName).build()) |
| .setStagingSessionToken(stagingSessionToken)) |
| .build()); |
| |
| try (FileInputStream fileInputStream = new FileInputStream(filePath)) { |
| byte[] buffer = new byte[DATA_1KB]; |
| int len; |
| while ((len = fileInputStream.read(buffer)) != -1) { |
| outputStreamObserver.onNext( |
| PutArtifactRequest.newBuilder() |
| .setData( |
| ArtifactChunk.newBuilder().setData(ByteString.copyFrom(buffer, 0, len)).build()) |
| .build()); |
| } |
| outputStreamObserver.onCompleted(); |
| complete.get(10, TimeUnit.SECONDS); |
| } |
| } |
| |
| private String commitManifest(String stagingSessionToken, List<ArtifactMetadata> artifacts) { |
| return stagingBlockingStub |
| .commitManifest( |
| CommitManifestRequest.newBuilder() |
| .setStagingSessionToken(stagingSessionToken) |
| .setManifest(Manifest.newBuilder().addAllArtifact(artifacts).build()) |
| .build()) |
| .getRetrievalToken(); |
| } |
| |
| @Test |
| public void generateStagingSessionTokenTest() throws Exception { |
| String basePath = stagingDir.toAbsolutePath().toString(); |
| String stagingToken = |
| BeamFileSystemArtifactStagingService.generateStagingSessionToken("abc123", basePath); |
| Assert.assertEquals( |
| "{\"sessionId\":\"abc123\",\"basePath\":\"" + basePath + "\"}", stagingToken); |
| } |
| |
| void checkCleanup(String stagingSessionToken, String stagingSession) throws Exception { |
| Assert.assertTrue( |
| Files.exists(Paths.get(stagingDir.toAbsolutePath().toString(), stagingSession))); |
| stagingService.removeArtifacts(stagingSessionToken); |
| Assert.assertFalse( |
| Files.exists(Paths.get(stagingDir.toAbsolutePath().toString(), stagingSession))); |
| } |
| |
| @Test |
| public void putArtifactsSingleSmallFileTest() throws Exception { |
| String fileName = "file1"; |
| String stagingSession = "123"; |
| String stagingSessionToken = |
| BeamFileSystemArtifactStagingService.generateStagingSessionToken( |
| stagingSession, stagingDir.toUri().getPath()); |
| Path srcFilePath = Paths.get(originalDir.toString(), fileName).toAbsolutePath(); |
| Files.write(srcFilePath, "some_test".getBytes(StandardCharsets.UTF_8)); |
| putArtifact(stagingSessionToken, srcFilePath.toString(), fileName); |
| String stagingToken = |
| commitManifest( |
| stagingSessionToken, |
| Collections.singletonList(ArtifactMetadata.newBuilder().setName(fileName).build())); |
| Assert.assertEquals( |
| Paths.get( |
| stagingDir.toAbsolutePath().toString(), |
| stagingSession, |
| BeamFileSystemArtifactStagingService.MANIFEST), |
| Paths.get(stagingToken)); |
| assertFiles(Collections.singleton(fileName), stagingToken); |
| checkCleanup(stagingSessionToken, stagingSession); |
| } |
| |
| @Test |
| public void putArtifactsMultipleFilesTest() throws Exception { |
| String stagingSession = "123"; |
| Map<String, Integer> files = |
| ImmutableMap.<String, Integer>builder() |
| .put("file5cb", (DATA_1KB / 2) /*500b*/) |
| .put("file1kb", DATA_1KB /*1 kb*/) |
| .put("file15cb", (DATA_1KB * 3) / 2 /*1.5 kb*/) |
| .put("nested/file1kb", DATA_1KB /*1 kb*/) |
| .put("file10kb", 10 * DATA_1KB /*10 kb*/) |
| .put("file100kb", 100 * DATA_1KB /*100 kb*/) |
| .build(); |
| Map<String, String> hashes = Maps.newHashMap(); |
| |
| final String text = "abcdefghinklmop\n"; |
| files.forEach( |
| (fileName, size) -> { |
| Path filePath = Paths.get(originalDir.toString(), fileName).toAbsolutePath(); |
| try { |
| Files.createDirectories(filePath.getParent()); |
| byte[] contents = |
| Strings.repeat( |
| text, Double.valueOf(Math.ceil(size * 1.0 / text.length())).intValue()) |
| .getBytes(StandardCharsets.UTF_8); |
| Files.write(filePath, contents); |
| hashes.put(fileName, Hashing.sha256().hashBytes(contents).toString()); |
| } catch (IOException ignored) { |
| } |
| }); |
| String stagingSessionToken = |
| BeamFileSystemArtifactStagingService.generateStagingSessionToken( |
| stagingSession, stagingDir.toUri().getPath()); |
| |
| List<ArtifactMetadata> metadata = new ArrayList<>(); |
| for (String fileName : files.keySet()) { |
| putArtifact( |
| stagingSessionToken, |
| Paths.get(originalDir.toString(), fileName).toAbsolutePath().toString(), |
| fileName); |
| metadata.add( |
| ArtifactMetadata.newBuilder().setName(fileName).setSha256(hashes.get(fileName)).build()); |
| } |
| |
| String retrievalToken = commitManifest(stagingSessionToken, metadata); |
| Assert.assertEquals( |
| Paths.get(stagingDir.toAbsolutePath().toString(), stagingSession, "MANIFEST").toString(), |
| retrievalToken); |
| assertFiles(files.keySet(), retrievalToken); |
| |
| checkCleanup(stagingSessionToken, stagingSession); |
| } |
| |
| @Test |
| public void putArtifactsMultipleFilesConcurrentlyTest() throws Exception { |
| String stagingSession = "123"; |
| Map<String, Integer> files = |
| ImmutableMap.<String, Integer>builder() |
| .put("file5cb", (DATA_1KB / 2) /*500b*/) |
| .put("file1kb", DATA_1KB /*1 kb*/) |
| .put("file15cb", (DATA_1KB * 3) / 2 /*1.5 kb*/) |
| .put("nested/file1kb", DATA_1KB /*1 kb*/) |
| .put("file10kb", 10 * DATA_1KB /*10 kb*/) |
| .put("file100kb", 100 * DATA_1KB /*100 kb*/) |
| .build(); |
| |
| final String text = "abcdefghinklmop\n"; |
| files.forEach( |
| (fileName, size) -> { |
| Path filePath = Paths.get(originalDir.toString(), fileName).toAbsolutePath(); |
| try { |
| Files.createDirectories(filePath.getParent()); |
| Files.write( |
| filePath, |
| Strings.repeat( |
| text, Double.valueOf(Math.ceil(size * 1.0 / text.length())).intValue()) |
| .getBytes(StandardCharsets.UTF_8)); |
| } catch (IOException ignored) { |
| } |
| }); |
| String stagingSessionToken = |
| BeamFileSystemArtifactStagingService.generateStagingSessionToken( |
| stagingSession, stagingDir.toUri().getPath()); |
| |
| List<ArtifactMetadata> metadata = Collections.synchronizedList(new ArrayList<>()); |
| ExecutorService executorService = Executors.newFixedThreadPool(8); |
| try { |
| for (String fileName : files.keySet()) { |
| executorService.execute( |
| () -> { |
| try { |
| putArtifact( |
| stagingSessionToken, |
| Paths.get(originalDir.toString(), fileName).toAbsolutePath().toString(), |
| fileName); |
| } catch (Exception e) { |
| Assert.fail(e.getMessage()); |
| } |
| metadata.add(ArtifactMetadata.newBuilder().setName(fileName).build()); |
| }); |
| } |
| } finally { |
| executorService.shutdown(); |
| executorService.awaitTermination(2, TimeUnit.SECONDS); |
| } |
| |
| String retrievalToken = commitManifest(stagingSessionToken, metadata); |
| Assert.assertEquals( |
| Paths.get(stagingDir.toAbsolutePath().toString(), stagingSession, "MANIFEST").toString(), |
| retrievalToken); |
| assertFiles(files.keySet(), retrievalToken); |
| |
| checkCleanup(stagingSessionToken, stagingSession); |
| } |
| |
| @Test |
| public void putArtifactsMultipleFilesConcurrentSessionsTest() throws Exception { |
| String stagingSession1 = "123"; |
| String stagingSession2 = "abc"; |
| Map<String, Integer> files1 = |
| ImmutableMap.<String, Integer>builder() |
| .put("file5cb", (DATA_1KB / 2) /*500b*/) |
| .put("file1kb", DATA_1KB /*1 kb*/) |
| .put("file15cb", (DATA_1KB * 3) / 2 /*1.5 kb*/) |
| .build(); |
| Map<String, Integer> files2 = |
| ImmutableMap.<String, Integer>builder() |
| .put("nested/file1kb", DATA_1KB /*1 kb*/) |
| .put("file10kb", 10 * DATA_1KB /*10 kb*/) |
| .put("file100kb", 100 * DATA_1KB /*100 kb*/) |
| .build(); |
| |
| final String text = "abcdefghinklmop\n"; |
| ImmutableMap.<String, Integer>builder() |
| .putAll(files1) |
| .putAll(files2) |
| .build() |
| .forEach( |
| (fileName, size) -> { |
| Path filePath = Paths.get(originalDir.toString(), fileName).toAbsolutePath(); |
| try { |
| Files.createDirectories(filePath.getParent()); |
| Files.write( |
| filePath, |
| Strings.repeat( |
| text, Double.valueOf(Math.ceil(size * 1.0 / text.length())).intValue()) |
| .getBytes(StandardCharsets.UTF_8)); |
| } catch (IOException ignored) { |
| } |
| }); |
| String stagingSessionToken1 = |
| BeamFileSystemArtifactStagingService.generateStagingSessionToken( |
| stagingSession1, stagingDir.toUri().getPath()); |
| String stagingSessionToken2 = |
| BeamFileSystemArtifactStagingService.generateStagingSessionToken( |
| stagingSession2, stagingDir.toUri().getPath()); |
| |
| List<ArtifactMetadata> metadata1 = Collections.synchronizedList(new ArrayList<>()); |
| List<ArtifactMetadata> metadata2 = Collections.synchronizedList(new ArrayList<>()); |
| ExecutorService executorService = Executors.newFixedThreadPool(8); |
| try { |
| Iterator<String> iterator1 = files1.keySet().iterator(); |
| Iterator<String> iterator2 = files2.keySet().iterator(); |
| while (iterator1.hasNext() && iterator2.hasNext()) { |
| String fileName1 = iterator1.next(); |
| String fileName2 = iterator2.next(); |
| executorService.execute( |
| () -> { |
| try { |
| putArtifact( |
| stagingSessionToken1, |
| Paths.get(originalDir.toString(), fileName1).toAbsolutePath().toString(), |
| fileName1); |
| putArtifact( |
| stagingSessionToken2, |
| Paths.get(originalDir.toString(), fileName2).toAbsolutePath().toString(), |
| fileName2); |
| } catch (Exception e) { |
| Assert.fail(e.getMessage()); |
| } |
| metadata1.add(ArtifactMetadata.newBuilder().setName(fileName1).build()); |
| metadata2.add(ArtifactMetadata.newBuilder().setName(fileName2).build()); |
| }); |
| } |
| } finally { |
| executorService.shutdown(); |
| executorService.awaitTermination(2, TimeUnit.SECONDS); |
| } |
| |
| String retrievalToken1 = commitManifest(stagingSessionToken1, metadata1); |
| String retrievalToken2 = commitManifest(stagingSessionToken2, metadata2); |
| Assert.assertEquals( |
| Paths.get(stagingDir.toAbsolutePath().toString(), stagingSession1, "MANIFEST").toString(), |
| retrievalToken1); |
| Assert.assertEquals( |
| Paths.get(stagingDir.toAbsolutePath().toString(), stagingSession2, "MANIFEST").toString(), |
| retrievalToken2); |
| assertFiles(files1.keySet(), retrievalToken1); |
| assertFiles(files2.keySet(), retrievalToken2); |
| |
| checkCleanup(stagingSessionToken1, stagingSession1); |
| checkCleanup(stagingSessionToken2, stagingSession2); |
| } |
| |
| private void assertFiles(Set<String> files, String retrievalToken) throws Exception { |
| ProxyManifest proxyManifest = |
| BeamFileSystemArtifactRetrievalService.loadManifest(retrievalToken); |
| GetManifestResponse retrievedManifest = |
| retrievalBlockingStub.getManifest( |
| GetManifestRequest.newBuilder().setRetrievalToken(retrievalToken).build()); |
| Assert.assertEquals( |
| "Manifest in proxy manifest file doesn't match the retrieved manifest", |
| proxyManifest.getManifest(), |
| retrievedManifest.getManifest()); |
| Assert.assertEquals( |
| "Files in locations does not match actual file list.", |
| files, |
| proxyManifest.getLocationList().stream() |
| .map(Location::getName) |
| .collect(Collectors.toSet())); |
| Assert.assertEquals( |
| "Duplicate file entries in locations.", files.size(), proxyManifest.getLocationCount()); |
| for (Location location : proxyManifest.getLocationList()) { |
| String expectedContent = |
| readFile( |
| Paths.get(originalDir.toString(), location.getName()).toAbsolutePath().toString()); |
| String actualStagedContent = readFile(location.getUri()); |
| Assert.assertEquals( |
| "Staged content doesn't match expected content for " + location.getName(), |
| expectedContent, |
| actualStagedContent); |
| String retrievedContent = retrieveArtifact(location.getName(), retrievalToken); |
| Assert.assertEquals( |
| "Retrieved content doesn't match expected content for " + location.getName(), |
| expectedContent, |
| retrievedContent); |
| } |
| } |
| |
| private static String readFile(String path) throws IOException { |
| try (InputStream stream = |
| Channels.newInputStream( |
| FileSystems.open(FileSystems.matchNewResource(path, false /* isDirectory */)))) { |
| return new String(ByteStreams.toByteArray(stream), StandardCharsets.UTF_8); |
| } |
| } |
| |
| private String retrieveArtifact(String name, String retrievalToken) |
| throws ExecutionException, InterruptedException { |
| CompletableFuture<ByteString> result = new CompletableFuture<>(); |
| retrievalStub.getArtifact( |
| GetArtifactRequest.newBuilder().setRetrievalToken(retrievalToken).setName(name).build(), |
| new StreamObserver<ArtifactChunk>() { |
| private ByteString data = ByteString.EMPTY; |
| |
| @Override |
| public void onNext(ArtifactChunk artifactChunk) { |
| data = data.concat(artifactChunk.getData()); |
| } |
| |
| @Override |
| public void onError(Throwable throwable) { |
| result.completeExceptionally(throwable); |
| } |
| |
| @Override |
| public void onCompleted() { |
| result.complete(data); |
| } |
| }); |
| return result.get().toStringUtf8(); |
| } |
| } |
