Add commands to preload files into S3

cmd/llama/main.go

@@ -38,6 +38,7 @@ func main() {
 	subcommands.Register(&XargsCommand{}, "")
 	subcommands.Register(&DaemonCommand{}, "")
 
+	subcommands.Register(&PreloadCommand{}, "internals")
 	subcommands.Register(&StoreCommand{}, "internals")
 	subcommands.Register(&GetCommand{}, "internals")
 	subcommands.Register(&trace.TraceCommand{}, "tracing")

cmd/llama/preload.go

@@ -0,0 +1,70 @@
+// Copyright 2020 Nelson Elhage
+//
+// 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 main
+
+import (
+	"context"
+	"flag"
+	"log"
+	"path/filepath"
+
+	"github.com/google/subcommands"
+	"github.com/nelhage/llama/cmd/internal/cli"
+	"github.com/nelhage/llama/daemon"
+	"github.com/nelhage/llama/daemon/server"
+)
+
+type PreloadCommand struct {
+	pattern string
+}
+
+func (*PreloadCommand) Name() string     { return "preload" }
+func (*PreloadCommand) Synopsis() string { return "Batch-upload files to the Llama store" }
+func (*PreloadCommand) Usage() string {
+	return `config PATH...
+`
+}
+
+func (c *PreloadCommand) SetFlags(flags *flag.FlagSet) {
+	flags.StringVar(&c.pattern, "pattern", `[.](c|cc|cxx|h|hpp|ipp)$`, "Regular expression defining files to preload")
+}
+
+func (c *PreloadCommand) Execute(ctx context.Context, flag *flag.FlagSet, _ ...interface{}) subcommands.ExitStatus {
+	client, err := server.DialWithAutostart(ctx, cli.SocketPath())
+	if err != nil {
+		log.Printf("Unable to connect to daemon: %s", err.Error())
+		return subcommands.ExitFailure
+	}
+	var args daemon.PreloadPathsArgs
+	for _, path := range flag.Args() {
+		abs, err := filepath.Abs(path)
+		if err != nil {
+			log.Fatalf("%q: %s", path, err.Error())
+			return subcommands.ExitFailure
+		}
+
+		args.Trees = append(args.Trees, daemon.PreloadTree{
+			Path:    abs,
+			Pattern: c.pattern,
+		})
+	}
+	out, err := client.PreloadPaths(&args)
+	if err != nil {
+		log.Printf("Preloading paths: %s", err.Error())
+		return subcommands.ExitFailure
+	}
+	log.Printf("Uploaded %d files to the Llama store", out.Preloaded)
+	return subcommands.ExitSuccess
+}

cmd/llamacc/main.go

@@ -82,11 +82,29 @@ func remap(local, wd string) files.Mapped {
 }
 
 func buildRemotePreprocess(ctx context.Context, client *daemon.Client, cfg *Config, comp *Compilation) error {
+	preload := make(chan struct{})
+	// buildRemoteInvoke will take some time to run `gcc` locally
+	// in order to discover dependencies. While it runs, we can at
+	// least start uploading the .c file we _know_ we will need to
+	// upload.
+	//
+	// During a build, we will likely eventually have uploaded all
+	// the header files, which means that this input file may be
+	// the _only_ source file we need to upload, and so getting a
+	// head-start on it can be a substantial performance
+	// improvement.
+	go func() {
+		defer close(preload)
+		_, span := tracing.StartSpan(ctx, "preload")
+		defer span.End()
+		client.PreloadPaths(&daemon.PreloadPathsArgs{Paths: []string{comp.Input}})
+	}()
 	args, err := buildRemoteInvoke(ctx, cfg, comp)
 	if err != nil {
 		return err
 	}
 	args.Trace = tracing.PropagationFromContext(ctx)
+	<-preload
 	out, err := client.InvokeWithFiles(args)
 	if err != nil {
 		return err

daemon/client.go

@@ -53,3 +53,9 @@ func (c *Client) TraceSpans(in *TraceSpansArgs) (*TraceSpansReply, error) {
 	err := c.conn.Call("Daemon.TraceSpans", in, &out)
 	return &out, err
 }
+
+func (c *Client) PreloadPaths(in *PreloadPathsArgs) (*PreloadPathsResult, error) {
+	var out PreloadPathsResult
+	err := c.conn.Call("Daemon.PreloadPaths", in, &out)
+	return &out, err
+}

daemon/server/methods.go

@@ -16,9 +16,13 @@ package server
 
 import (
 	"fmt"
+	"io/fs"
+	"io/ioutil"
 	"log"
 	"os"
 	"path"
+	"path/filepath"
+	"regexp"
 	"sync/atomic"
 	"time"
 
@@ -28,6 +32,7 @@ import (
 	"github.com/nelhage/llama/protocol/files"
 	"github.com/nelhage/llama/store"
 	"github.com/nelhage/llama/tracing"
+	"golang.org/x/sync/errgroup"
 )
 
 func (d *Daemon) Ping(in daemon.PingArgs, reply *daemon.PingReply) error {
@@ -216,3 +221,67 @@ func (d *Daemon) TraceSpans(in *daemon.TraceSpansArgs, out *daemon.TraceSpansRep
 	*out = daemon.TraceSpansReply{}
 	return nil
 }
+
+const preloadThreads = 32
+
+func (d *Daemon) PreloadPaths(in *daemon.PreloadPathsArgs, out *daemon.PreloadPathsResult) error {
+	grp, ctx := errgroup.WithContext(d.ctx)
+	paths := make(chan string)
+	var preloaded uint64
+	for i := 0; i < preloadThreads; i++ {
+		grp.Go(func() error {
+			for {
+				select {
+				case path, ok := <-paths:
+					if !ok {
+						return nil
+					}
+					data, err := ioutil.ReadFile(path)
+					if err != nil {
+						return err
+					}
+					if _, err := d.store.Store(ctx, data); err != nil {
+						return err
+					}
+					atomic.AddUint64(&preloaded, 1)
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+			}
+		})
+	}
+	grp.Go(func() error {
+		defer close(paths)
+		for _, path := range in.Paths {
+			paths <- path
+		}
+		for _, req := range in.Trees {
+			re, err := regexp.Compile(req.Pattern)
+			if err != nil {
+				return err
+			}
+			err = filepath.WalkDir(req.Path, func(path string, d fs.DirEntry, err error) error {
+				if err != nil {
+					return err
+				}
+				if d.IsDir() {
+					return nil
+				}
+				if re.MatchString(d.Name()) {
+					paths <- path
+				}
+				return nil
+			})
+			if err != nil {
+				return err
+			}
+		}
+		return nil
+	})
+	err := grp.Wait()
+	if err != nil {
+		return err
+	}
+	*out = daemon.PreloadPathsResult{Preloaded: preloaded}
+	return nil
+}

daemon/types.go

@@ -78,3 +78,19 @@ type TraceSpansArgs struct {
 }
 
 type TraceSpansReply struct{}
+
+type PreloadTree struct {
+	Path string
+
+	// A regular expression; will be matched against the file
+	// basename.
+	Pattern string
+}
+
+type PreloadPathsArgs struct {
+	Paths []string
+	Trees []PreloadTree
+}
+type PreloadPathsResult struct {
+	Preloaded uint64
+}