com.google.devtools.build.remote.RemoteWorker.java Source code

Java tutorial

Introduction

Here is the source code for com.google.devtools.build.remote.RemoteWorker.java

Source

// Copyright 2016 The Bazel 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
//
//    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 com.google.devtools.build.remote;

import com.google.common.collect.ImmutableList;
import com.google.devtools.build.lib.remote.CacheNotFoundException;
import com.google.devtools.build.lib.remote.ConcurrentMapActionCache;
import com.google.devtools.build.lib.remote.ConcurrentMapFactory;
import com.google.devtools.build.lib.remote.ContentDigests;
import com.google.devtools.build.lib.remote.ExecuteServiceGrpc.ExecuteServiceImplBase;
import com.google.devtools.build.lib.remote.RemoteOptions;
import com.google.devtools.build.lib.remote.RemoteProtocol;
import com.google.devtools.build.lib.remote.RemoteProtocol.Action;
import com.google.devtools.build.lib.remote.RemoteProtocol.ActionResult;
import com.google.devtools.build.lib.remote.RemoteProtocol.CasStatus;
import com.google.devtools.build.lib.remote.RemoteProtocol.Command.EnvironmentEntry;
import com.google.devtools.build.lib.remote.RemoteProtocol.ContentDigest;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecuteReply;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecuteRequest;
import com.google.devtools.build.lib.remote.RemoteProtocol.ExecutionStatus;
import com.google.devtools.build.lib.shell.AbnormalTerminationException;
import com.google.devtools.build.lib.shell.Command;
import com.google.devtools.build.lib.shell.CommandException;
import com.google.devtools.build.lib.util.OS;
import com.google.devtools.build.lib.util.ProcessUtils;
import com.google.devtools.build.lib.vfs.FileSystem;
import com.google.devtools.build.lib.vfs.FileSystemUtils;
import com.google.devtools.build.lib.vfs.JavaIoFileSystem;
import com.google.devtools.build.lib.vfs.Path;
import com.google.devtools.build.lib.vfs.UnixFileSystem;
import com.google.devtools.common.options.OptionsParser;
import io.grpc.Server;
import io.grpc.ServerBuilder;
import io.grpc.stub.StreamObserver;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.file.FileAlreadyExistsException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * Implements a remote worker that accepts work items as protobufs. The server implementation is
 * based on grpc.
 */
public class RemoteWorker extends ExecuteServiceImplBase {
    private static final Logger LOG = Logger.getLogger(RemoteWorker.class.getName());
    private static final boolean LOG_FINER = LOG.isLoggable(Level.FINER);
    private final Path workPath;
    private final RemoteOptions remoteOptions;
    private final RemoteWorkerOptions options;
    private final ConcurrentMapActionCache cache;

    public RemoteWorker(Path workPath, RemoteOptions remoteOptions, RemoteWorkerOptions options,
            ConcurrentMapActionCache cache) {
        this.workPath = workPath;
        this.remoteOptions = remoteOptions;
        this.options = options;
        this.cache = cache;
    }

    private Map<String, String> getEnvironmentVariables(RemoteProtocol.Command command) {
        HashMap<String, String> result = new HashMap<>();
        for (EnvironmentEntry entry : command.getEnvironmentList()) {
            result.put(entry.getVariable(), entry.getValue());
        }
        return result;
    }

    public ExecuteReply execute(Action action, Path execRoot) throws IOException, InterruptedException {
        ByteArrayOutputStream stdout = new ByteArrayOutputStream();
        ByteArrayOutputStream stderr = new ByteArrayOutputStream();
        try {
            RemoteProtocol.Command command = RemoteProtocol.Command
                    .parseFrom(cache.downloadBlob(action.getCommandDigest()));
            cache.downloadTree(action.getInputRootDigest(), execRoot);

            List<Path> outputs = new ArrayList<>(action.getOutputPathList().size());
            for (String output : action.getOutputPathList()) {
                Path file = execRoot.getRelative(output);
                if (file.exists()) {
                    throw new FileAlreadyExistsException("Output file already exists: " + file);
                }
                FileSystemUtils.createDirectoryAndParents(file.getParentDirectory());
                outputs.add(file);
            }

            // TODO(olaola): time out after specified server-side deadline.
            Command cmd = new Command(command.getArgvList().toArray(new String[] {}),
                    getEnvironmentVariables(command), new File(execRoot.getPathString()));
            cmd.execute(Command.NO_INPUT, Command.NO_OBSERVER, stdout, stderr, true);

            // Execute throws a CommandException on non-zero return values, so action has succeeded.
            ImmutableList<ContentDigest> outErrDigests = cache
                    .uploadBlobs(ImmutableList.of(stdout.toByteArray(), stderr.toByteArray()));
            ActionResult.Builder result = ActionResult.newBuilder().setReturnCode(0)
                    .setStdoutDigest(outErrDigests.get(0)).setStderrDigest(outErrDigests.get(1));
            cache.uploadAllResults(execRoot, outputs, result);
            cache.setCachedActionResult(ContentDigests.computeActionKey(action), result.build());
            return ExecuteReply.newBuilder().setResult(result)
                    .setStatus(ExecutionStatus.newBuilder().setExecuted(true).setSucceeded(true)).build();
        } catch (CommandException e) {
            ImmutableList<ContentDigest> outErrDigests = cache
                    .uploadBlobs(ImmutableList.of(stdout.toByteArray(), stderr.toByteArray()));
            final int returnCode = e instanceof AbnormalTerminationException
                    ? ((AbnormalTerminationException) e).getResult().getTerminationStatus().getExitCode()
                    : -1;
            return ExecuteReply.newBuilder()
                    .setResult(ActionResult.newBuilder().setReturnCode(returnCode)
                            .setStdoutDigest(outErrDigests.get(0)).setStderrDigest(outErrDigests.get(1)))
                    .setStatus(ExecutionStatus.newBuilder().setExecuted(true).setSucceeded(false)
                            .setError(ExecutionStatus.ErrorCode.EXEC_FAILED).setErrorDetail(e.toString()))
                    .build();
        } catch (CacheNotFoundException e) {
            LOG.warning("Cache miss on " + ContentDigests.toString(e.getMissingDigest()));
            return ExecuteReply.newBuilder()
                    .setCasError(CasStatus.newBuilder().setSucceeded(false).addMissingDigest(e.getMissingDigest())
                            .setError(CasStatus.ErrorCode.MISSING_DIGEST).setErrorDetail(e.toString()))
                    .setStatus(ExecutionStatus.newBuilder().setExecuted(false).setSucceeded(false)
                            .setError(e.getMissingDigest() == action.getCommandDigest()
                                    ? ExecutionStatus.ErrorCode.MISSING_COMMAND
                                    : ExecutionStatus.ErrorCode.MISSING_INPUT)
                            .setErrorDetail(e.toString()))
                    .build();
        }
    }

    @Override
    public void execute(ExecuteRequest request, StreamObserver<ExecuteReply> responseObserver) {
        Path tempRoot = workPath.getRelative("build-" + UUID.randomUUID().toString());
        try {
            tempRoot.createDirectory();
            if (LOG_FINER) {
                LOG.fine("Work received has " + request.getTotalInputFileCount() + " input files and "
                        + request.getAction().getOutputPathCount() + " output files.");
            }
            ExecuteReply reply = execute(request.getAction(), tempRoot);
            responseObserver.onNext(reply);
            if (options.debug) {
                if (!reply.getStatus().getSucceeded()) {
                    LOG.warning("Work failed. Request: " + request.toString() + ".");
                } else if (LOG_FINER) {
                    LOG.fine("Work completed.");
                }
            }
            if (!options.debug) {
                FileSystemUtils.deleteTree(tempRoot);
            } else {
                LOG.warning("Preserving work directory " + tempRoot.toString() + ".");
            }
        } catch (IOException | InterruptedException e) {
            ExecuteReply.Builder reply = ExecuteReply.newBuilder();
            reply.getStatusBuilder().setSucceeded(false).setErrorDetail(e.toString());
            responseObserver.onNext(reply.build());
            if (e instanceof InterruptedException) {
                Thread.currentThread().interrupt();
            }
        } finally {
            responseObserver.onCompleted();
        }
    }

    public static void main(String[] args) throws Exception {
        OptionsParser parser = OptionsParser.newOptionsParser(RemoteOptions.class, RemoteWorkerOptions.class);
        parser.parseAndExitUponError(args);
        RemoteOptions remoteOptions = parser.getOptions(RemoteOptions.class);
        RemoteWorkerOptions remoteWorkerOptions = parser.getOptions(RemoteWorkerOptions.class);

        if (remoteWorkerOptions.workPath == null) {
            printUsage(parser);
            return;
        }

        System.out.println("*** Starting Hazelcast server.");
        ConcurrentMapActionCache cache = new ConcurrentMapActionCache(
                ConcurrentMapFactory.createHazelcast(remoteOptions));

        System.out.println("*** Starting grpc server on all locally bound IPs on port "
                + remoteWorkerOptions.listenPort + ".");
        Path workPath = getFileSystem().getPath(remoteWorkerOptions.workPath);
        FileSystemUtils.createDirectoryAndParents(workPath);
        RemoteWorker worker = new RemoteWorker(workPath, remoteOptions, remoteWorkerOptions, cache);
        final Server server = ServerBuilder.forPort(remoteWorkerOptions.listenPort).addService(worker).build();
        server.start();

        final Path pidFile;
        if (remoteWorkerOptions.pidFile != null) {
            pidFile = getFileSystem().getPath(remoteWorkerOptions.pidFile);
            PrintWriter writer = new PrintWriter(pidFile.getOutputStream());
            writer.append(Integer.toString(ProcessUtils.getpid()));
            writer.append("\n");
            writer.close();
        } else {
            pidFile = null;
        }

        Runtime.getRuntime().addShutdownHook(new Thread() {
            @Override
            public void run() {
                System.err.println("*** Shutting down grpc server.");
                server.shutdown();
                if (pidFile != null) {
                    try {
                        pidFile.delete();
                    } catch (IOException e) {
                        System.err.println("Cannot remove pid file: " + pidFile.toString());
                    }
                }
                System.err.println("*** Server shut down.");
            }
        });
        server.awaitTermination();
    }

    public static void printUsage(OptionsParser parser) {
        System.out.println("Usage: remote_worker \n\n" + "Starts a worker that runs a RPC service.");
        System.out.println(
                parser.describeOptions(Collections.<String, String>emptyMap(), OptionsParser.HelpVerbosity.LONG));
    }

    static FileSystem getFileSystem() {
        return OS.getCurrent() == OS.WINDOWS ? new JavaIoFileSystem() : new UnixFileSystem();
    }
}