WIP: bazel-bes uploads via channels etc

Author: pda

internal/bes/bes.go

@@ -8,47 +8,21 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"net"
 	"sort"
 
 	slog "github.com/buildkite/test-engine-client/internal/bes/quietslog"
 
 	"google.golang.org/genproto/googleapis/devtools/build/v1"
-	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/types/known/emptypb"
 )
 
-var host = "127.0.0.1"
-var port = 60242 // 0 for OS-allocated
-
 type BuildEventServer struct {
 	handler *BuildEventHandler
 }
 
-func Listen() error {
-	addr := fmt.Sprintf("%s:%d", host, port)
-	listener, err := net.Listen("tcp", addr)
-	if err != nil {
-		return fmt.Errorf("listening on %s: %w", addr, err)
-	}
-	fmt.Println("Bazel BES listener: grpc://" + listener.Addr().String())
-
-	opts := []grpc.ServerOption{}
-	grpcServer := grpc.NewServer(opts...)
-
-	build.RegisterPublishBuildEventServer(grpcServer, newServer())
-	grpcServer.Serve(listener)
-
-	return nil
-}
-
-func newServer() build.PublishBuildEventServer {
-	return BuildEventServer{}
-}
-
 // PublishLifecycleEvent is copied verbatim from:
 // https://github.com/buildbarn/bb-portal/blob/abb76f0a9324cf4f9d5da44b53804a8d9a0a2155/internal/api/grpc/bes/server.go
 //

internal/bes/bes_test.go

@@ -0,0 +1,14 @@
+package bes
+
+import "testing"
+
+func TestPathFromURI(t *testing.T) {
+	path, err := pathFromURI("file:///hello/world.txt")
+	if err != nil {
+		t.Errorf("pathFromURI error: %v", err)
+	}
+
+	if want := "/hello/world.txt"; want != path {
+		t.Errorf("wanted %v got %v", want, path)
+	}
+}

internal/bes/channel.go

@@ -2,7 +2,9 @@ package bes
 
 import (
 	"context"
+	"fmt"
 	"log/slog"
+	"net/url"
 	"time"
 
 	"github.com/buildbarn/bb-portal/third_party/bazel/gen/bes"
@@ -24,8 +26,9 @@ type BuildEventChannel interface {
 }
 
 type buildEventChannel struct {
-	ctx      context.Context
-	streamID *build.StreamId
+	ctx       context.Context
+	streamID  *build.StreamId
+	filenames chan<- string
 }
 
 // HandleBuildEvent implements BuildEventChannel.HandleBuildEvent.
@@ -45,7 +48,12 @@ func (c *buildEventChannel) HandleBuildEvent(event *build.BuildEvent) error {
 		files := []string{}
 		for _, x := range r.GetTestActionOutput() {
 			if x.GetName() == "test.xml" {
-				files = append(files, x.GetUri())
+				path, err := pathFromURI(x.GetUri())
+				if err != nil {
+					return err // maybe just a log a warning?
+				}
+				files = append(files, path)
+				c.filenames <- path
 			}
 		}
 		slog.Info("TestResult",
@@ -59,13 +67,24 @@ func (c *buildEventChannel) HandleBuildEvent(event *build.BuildEvent) error {
 	return nil
 }
 
+func pathFromURI(uri string) (string, error) {
+	url, err := url.Parse(uri)
+	if err != nil {
+		return "", err
+	}
+	if url.Scheme != "file" {
+		return "", fmt.Errorf("expected file://..., got %v://...", url.Scheme)
+	}
+	return url.Path, nil
+}
+
 // Finalize implements BuildEventChannel.Finalize.
 func (c *buildEventChannel) Finalize() error {
 	// defer the ctx so its not reaped when the client closes the connection
 	ctx, cancel := context.WithTimeout(context.Background(), time.Hour*24)
 	defer cancel()
 
-	slog.Debug("finalizing build event channel")
+	slog.Info("finalizing build event channel")
 	_ = ctx
 	// TODO: finalize anything that needs finalizing?
 

internal/bes/handler.go

@@ -13,11 +13,14 @@ import (
 // For each incoming stream, and BuildEventChannel is created, which handles that stream.
 // BuildEventHandler is responsible for managing the things that are common to these event streams.
 type BuildEventHandler struct {
+	filenames chan<- string
 }
 
 // NewBuildEventHandler constructs a new BuildEventHandler
-func NewBuildEventHandler() *BuildEventHandler {
-	return &BuildEventHandler{}
+func NewBuildEventHandler(filenames chan<- string) *BuildEventHandler {
+	return &BuildEventHandler{
+		filenames: filenames,
+	}
 }
 
 // CreateEventChannel creates a new BuildEventChannel
@@ -29,7 +32,8 @@ func (h *BuildEventHandler) CreateEventChannel(ctx context.Context, initialEvent
 	}
 
 	return &buildEventChannel{
-		ctx:      ctx,
-		streamID: initialEvent.StreamId,
+		ctx:       ctx,
+		streamID:  initialEvent.StreamId,
+		filenames: h.filenames,
 	}
 }

internal/bes/listen.go

@@ -0,0 +1,101 @@
+package bes
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"log/slog"
+	"net"
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/buildkite/test-engine-client/internal/env"
+	"github.com/buildkite/test-engine-client/internal/upload"
+	"google.golang.org/genproto/googleapis/devtools/build/v1"
+	"google.golang.org/grpc"
+)
+
+func ListenCLI(argv []string, env env.Env) error {
+	flags := flag.NewFlagSet("bktec bazel listen", flag.ExitOnError)
+	portFlag := flags.Int("port", 0, "gRPC port to listen")
+	listenHostFlag := flags.String("listen-host", "127.0.0.1", "gRPC host to listen")
+	debugFlag := flags.Bool("debug", false, "debug logging")
+	flags.Parse(argv)
+
+	if *debugFlag {
+		slog.SetLogLoggerLevel(slog.LevelDebug)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// a channel to propagate OS signals
+	signals := make(chan os.Signal, 1)
+	signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT)
+
+	// configure uploader
+	cfg, err := upload.ConfigFromEnv(env)
+	if err != nil {
+		return fmt.Errorf("uploader configuration: %w", err)
+	}
+	runEnv, err := upload.RunEnvFromEnv(env)
+	if err != nil {
+		return fmt.Errorf("uploader run_env configuration: %w", err)
+	}
+	uploader := NewUploader(cfg, runEnv, "junit")
+	go uploader.Start(ctx)
+
+	// configure gRPC Bazel BES server
+	addr := fmt.Sprintf("%s:%d", *listenHostFlag, *portFlag)
+	listener, err := net.Listen("tcp", addr)
+	if err != nil {
+		return fmt.Errorf("listening on %s: %w", addr, err)
+	}
+	opts := []grpc.ServerOption{}
+	srv := grpc.NewServer(opts...)
+	build.RegisterPublishBuildEventServer(srv, BuildEventServer{
+		handler: &BuildEventHandler{
+			filenames: uploader.Filenames,
+		},
+	})
+	slog.Info("Bazel BES listener", "addr", "grpc://"+listener.Addr().String())
+	go serve(srv, listener)
+
+	// main loop
+	run := true
+	sigCount := 0
+	for run {
+		select {
+		case url, ok := <-uploader.Responses:
+			if !ok {
+				slog.Debug("Response channel closed")
+				run = false
+				continue
+			}
+			slog.Info("Uploaded", "url", url)
+		case err := <-uploader.Errs:
+			slog.Error("Upload error", "error", err)
+		case sig := <-signals:
+			sigCount++
+			srv.Stop()
+			if sigCount == 1 {
+				slog.Info("Stopping (again to force)...", "signal", sig)
+				uploader.Stop()
+			} else {
+				slog.Info("Stopping forcefully...", "signal", sig)
+				cancel()
+			}
+		}
+	}
+
+	slog.Debug("done")
+	return nil
+}
+
+func serve(s *grpc.Server, listener net.Listener) {
+	err := s.Serve(listener)
+	if err != nil {
+		slog.Error("gRPC server error", "err", err)
+	}
+}

internal/bes/uploader.go

@@ -0,0 +1,76 @@
+package bes
+
+import (
+	"context"
+	"log/slog"
+
+	"github.com/buildkite/test-engine-client/internal/upload"
+)
+
+type Uploader struct {
+	Config    upload.Config
+	RunEnv    upload.RunEnvMap
+	Format    string
+	Filenames chan string
+	Responses chan string
+	Errs      chan error
+
+	stopping bool
+}
+
+func NewUploader(cfg upload.Config, runEnv upload.RunEnvMap, format string) *Uploader {
+	// a channel to pass filenames from BES server to uploader
+	filenames := make(chan string, 1024)
+
+	// a channel to receive response upload URLs
+	responses := make(chan string)
+
+	// a channel to receive errors from the uploader
+	errs := make(chan error)
+
+	return &Uploader{
+		Config:    cfg,
+		RunEnv:    runEnv,
+		Format:    format,
+		Filenames: filenames,
+		Responses: responses,
+		Errs:      errs,
+	}
+}
+
+func (u *Uploader) Start(ctx context.Context) {
+	for filename := range u.Filenames {
+		if ctx.Err() != nil {
+			slog.Debug("Uploader context canceled")
+			break
+		}
+		resp, err := u.UploadFile(ctx, filename)
+		if err != nil {
+			u.Errs <- err
+			continue
+		}
+		u.Responses <- resp["upload_url"]
+	}
+	slog.Debug("Uploader finished")
+	close(u.Responses)
+}
+
+// Stop closes the Filenames channel; filenames already buffered on the channel
+// will be uploaded before finishing.
+func (u *Uploader) Stop() {
+	if u.stopping {
+		slog.Warn("Uploader GracefulStop: already stopping")
+		return
+	}
+	slog.Debug("Uploader GracefulStop")
+	u.stopping = true
+	close(u.Filenames)
+}
+
+func (u *Uploader) UploadFile(ctx context.Context, filename string) (map[string]string, error) {
+	resp, err := upload.Upload(ctx, u.Config, u.RunEnv, u.Format, filename)
+	if err != nil {
+		return nil, err
+	}
+	return resp, nil
+}

main.go

@@ -72,7 +72,7 @@ func main() {
 		}
 		os.Exit(0)
 	} else if flag.Arg(0) == "bazel" && flag.Arg(1) == "listen" {
-		if err := bes.Listen(); err != nil {
+		if err := bes.ListenCLI(os.Args[3:], env); err != nil {
 			log.Fatal(err)
 		}
 		os.Exit(0)