From 879096ea97e727fa06d7db84522e457b199babc3 Mon Sep 17 00:00:00 2001 From: Sairaj Pokale Date: Fri, 28 Aug 2026 14:17:58 -0700 Subject: [PATCH 1/2] benchmarking: move glutton's core into internal/benchmarking/glutton cmd/benchmarking/glutton/main.go held the entire workload: the service, its otel instruments, the gossip loop, and both transports. None of it was reachable outside package main, which left internal/benchmarking/glutton holding nothing but a fake of a server implemented elsewhere. Move everything except process lifecycle into the internal package, matching cmd/benchmarking/boomer-glutton. main.go keeps flag parsing, telemetry boot, the listener, and the shutdown defers. The --mode switch moves out too, since it selects a transport rather than managing the process, and now returns an error for an unknown mode instead of calling serverboot.Fatal. The destination is top-level internal/ because both packages that want to share glutton's route contract (the fake, and boomer's glutton user class) sit outside cmd/benchmarking/glutton/. Renames: gluttonService to Service, newGluttonService to New, newServer to NewServer, and the mode switch to Handler. sizes.go moves along with its only caller. No behavior change. Signed-off-by: Sairaj Pokale --- cmd/benchmarking/glutton/main.go | 694 +---------------- internal/benchmarking/glutton/glutton.go | 726 ++++++++++++++++++ .../benchmarking/glutton/glutton_test.go | 78 +- .../benchmarking/glutton/ram_test.go | 44 +- .../benchmarking/glutton/sizes.go | 2 +- .../benchmarking/glutton/sizes_test.go | 16 +- 6 files changed, 802 insertions(+), 758 deletions(-) create mode 100644 internal/benchmarking/glutton/glutton.go rename cmd/benchmarking/glutton/main_test.go => internal/benchmarking/glutton/glutton_test.go (85%) rename {cmd => internal}/benchmarking/glutton/ram_test.go (78%) rename {cmd => internal}/benchmarking/glutton/sizes.go (98%) rename {cmd => internal}/benchmarking/glutton/sizes_test.go (77%) diff --git a/cmd/benchmarking/glutton/main.go b/cmd/benchmarking/glutton/main.go index 5488402745..126270054d 100644 --- a/cmd/benchmarking/glutton/main.go +++ b/cmd/benchmarking/glutton/main.go @@ -19,44 +19,18 @@ package main import ( "context" - "crypto/rand" - "crypto/sha256" - "errors" "fmt" - "io" "log/slog" "net" - "net/http" "os" - "path/filepath" - "regexp" - "strconv" - "strings" - "sync" - "time" - "github.com/google/uuid" "github.com/spf13/pflag" - "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" - "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" - "go.opentelemetry.io/otel" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/metric" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials/insecure" - "google.golang.org/grpc/reflection" - "google.golang.org/grpc/status" - "google.golang.org/protobuf/proto" - "github.com/agent-substrate/substrate/internal/ateinterceptors" - "github.com/agent-substrate/substrate/internal/proto/glutton" + "github.com/agent-substrate/substrate/internal/benchmarking/glutton" "github.com/agent-substrate/substrate/internal/serverboot" "github.com/agent-substrate/substrate/internal/version" ) -const meterName = "glutton" - var ( listenAddr = pflag.String("grpc-listen-addr", ":8080", "Address and port the server should listen on (name kept for back-compat; serves whatever --mode picks).") metricsListenAddr = pflag.String("metrics-listen-addr", ":9090", "Address and port the Prometheus metrics server should listen on.") @@ -99,7 +73,7 @@ func main() { serverboot.Fatal(ctx, "Failed to create data directory", fmt.Errorf("%s: %w", *dataDir, err)) } - svc, err := newGluttonService(*dataDir) + svc, err := glutton.New(*dataDir) if err != nil { serverboot.Fatal(ctx, "Failed to construct glutton service", err) } @@ -122,667 +96,11 @@ func main() { slog.String("mode", *mode), ) - var handler http.Handler - switch *mode { - case "grpc": - srv := grpc.NewServer( - grpc.StatsHandler(otelgrpc.NewServerHandler()), - ) - glutton.RegisterGluttonServer(srv, svc) - reflection.Register(srv) - // The readiness probe is an HTTP GET, so gRPC mode serves it next to - // the gRPC handler on the same listener. - handler = splitGRPC(srv, readyzMux()) - case "http": - // otelhttp at the mux level + per-handler span follows - // docs/dev/best-practices/tracing.md: extract incoming context, - // then name the span after the operation in each handler. - handler = otelhttp.NewHandler(newMux(svc), "/") - default: - serverboot.Fatal(ctx, "Invalid --mode", fmt.Errorf("must be grpc or http: %q", *mode)) - } - if err := newServer(handler).Serve(lis); err != nil { - serverboot.Fatal(ctx, "Failed to serve", err) - } -} - -// newServer enables unencrypted HTTP/2 so gRPC works on the plaintext -// listener, alongside HTTP/1.1 for the readyz probe. -func newServer(handler http.Handler) *http.Server { - protocols := new(http.Protocols) - protocols.SetHTTP1(true) - protocols.SetUnencryptedHTTP2(true) - return &http.Server{Handler: handler, Protocols: protocols} -} - -// splitGRPC serves gRPC and plain HTTP on one listener: requests with a -// gRPC content-type go to grpcSrv, everything else to rest. All glutton -// RPCs are unary, which is what grpc.Server.ServeHTTP supports. -func splitGRPC(grpcSrv, rest http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") { - grpcSrv.ServeHTTP(w, r) - return - } - rest.ServeHTTP(w, r) - }) -} - -// readyzMux serves the readiness probe both modes need: ateom blocks -// RestoreWorkload until /readyz returns 200, so ResumeActor cannot report -// success before this listener is reachable. -func readyzMux() *http.ServeMux { - mux := http.NewServeMux() - mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - }) - return mux -} - -// newMux builds the HTTP-mode route table on top of the readiness probe. -func newMux(svc *gluttonService) *http.ServeMux { - mux := readyzMux() - mux.HandleFunc("/ping", protoRoute("Ping", svc.Ping)) - mux.HandleFunc("/writedisk", protoRoute("WriteDisk", svc.WriteDisk)) - mux.HandleFunc("/readdisk", protoRoute("ReadDisk", svc.ReadDisk)) - mux.HandleFunc("/writeram", protoRoute("WriteRAM", svc.WriteRAM)) - mux.HandleFunc("/readram", protoRoute("ReadRAM", svc.ReadRAM)) - return mux -} - -// protoRoute wraps a protobuf handler with POST-only routing, protobuf -// unmarshaling, status code mapping, and server-timing headers. -func protoRoute[Req any, Resp proto.Message, PtrReq interface { - *Req - proto.Message -}](spanName string, handler func(context.Context, PtrReq) (Resp, error)) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - start := time.Now() - if r.Method != http.MethodPost { - http.Error(w, "method not allowed", http.StatusMethodNotAllowed) - return - } - body, err := io.ReadAll(r.Body) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - var req Req - ptrReq := PtrReq(&req) - if err := proto.Unmarshal(body, ptrReq); err != nil { - http.Error(w, "unmarshal: "+err.Error(), http.StatusBadRequest) - return - } - ctx, span := otel.Tracer("glutton").Start(r.Context(), spanName) - defer span.End() - resp, err := handler(ctx, ptrReq) - if err != nil { - if st, ok := status.FromError(err); ok { - switch st.Code() { - case codes.InvalidArgument: - http.Error(w, st.Message(), http.StatusBadRequest) - case codes.NotFound: - http.Error(w, st.Message(), http.StatusNotFound) - default: - http.Error(w, st.Message(), http.StatusInternalServerError) - } - } else { - http.Error(w, err.Error(), http.StatusInternalServerError) - } - return - } - out, err := proto.Marshal(resp) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - // Glutton does not run ateinterceptors, so without this the serve path has no - // server-side timing at all. Mirrors the control-plane gRPC trailer so boomer's - // elapsedFromMD logic (source=server) works identically over HTTP. - w.Header().Set(ateinterceptors.ServerElapsedTrailer, - strconv.FormatInt(time.Since(start).Microseconds(), 10)) - w.Header().Set("Content-Type", "application/x-protobuf") - _, _ = w.Write(out) - } -} - -// diskKeyRE rejects anything that could escape the data dir or hit a -// hidden file: only alphanumerics, underscore, and dash are permitted. -var diskKeyRE = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`) - -type gluttonService struct { - glutton.UnimplementedGluttonServer - - dataDir string - - // TODO: split this into per-resource locks (ram, fds, peers). A single - // global mutex serializes unrelated operations across all three. - mu sync.Mutex - ram map[string][]byte - // ramCursor is each array's next WRITE_MODE_OVERWRITE_ROTATE offset. - // Absent means 0; invalidated whenever the array is reallocated. - ramCursor map[string]int - fds []*os.File - peers map[string]*peerGossip - - ramWriteBytes metric.Int64Counter - ramReadBytes metric.Int64Counter - diskWriteBytes metric.Int64Counter - diskReadBytes metric.Int64Counter - pingsReceived metric.Int64Counter - gossipSent metric.Int64Counter - gossipLatency metric.Float64Histogram -} - -type peerGossip struct { - host string - delayMs int32 - cancel context.CancelFunc - done chan struct{} -} - -func newGluttonService(dir string) (*gluttonService, error) { - s := &gluttonService{ - dataDir: dir, - ram: make(map[string][]byte), - ramCursor: make(map[string]int), - peers: make(map[string]*peerGossip), - } - - m := otel.Meter(meterName) - - var err error - s.ramWriteBytes, err = m.Int64Counter( - "glutton.ram.write.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes written to RAM via WriteRAM over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.ram.write.bytes counter: %w", err) - } - s.ramReadBytes, err = m.Int64Counter( - "glutton.ram.read.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes walked by ReadRAM over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.ram.read.bytes counter: %w", err) - } - s.diskWriteBytes, err = m.Int64Counter( - "glutton.disk.write.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes written to disk via WriteDisk over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.disk.write.bytes counter: %w", err) - } - s.diskReadBytes, err = m.Int64Counter( - "glutton.disk.read.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes read from disk via ReadDisk over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.disk.read.bytes counter: %w", err) - } - s.pingsReceived, err = m.Int64Counter( - "glutton.ping.requests", - metric.WithDescription("Number of Ping requests received."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.ping.requests counter: %w", err) - } - s.gossipSent, err = m.Int64Counter( - "glutton.gossip.requests.sent", - metric.WithDescription("Number of gossip Ping requests sent per peer."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.gossip.requests.sent counter: %w", err) - } - s.gossipLatency, err = m.Float64Histogram( - "glutton.gossip.latency", - metric.WithUnit("s"), - metric.WithDescription("Latency of gossip Ping requests per peer."), - metric.WithExplicitBucketBoundaries( - 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, - ), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.gossip.latency histogram: %w", err) - } - - fdsOpen, err := m.Int64ObservableGauge( - "glutton.fds.open", - metric.WithDescription("File descriptors currently held open by OpenFD."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.fds.open gauge: %w", err) - } - peerDelay, err := m.Int64ObservableGauge( - "glutton.gossip.delay", - metric.WithUnit("ms"), - metric.WithDescription("Configured gossip delay per peer."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.gossip.delay gauge: %w", err) - } - - if _, err := m.RegisterCallback(func(_ context.Context, o metric.Observer) error { - s.mu.Lock() - defer s.mu.Unlock() - o.ObserveInt64(fdsOpen, int64(len(s.fds))) - for host, p := range s.peers { - o.ObserveInt64(peerDelay, int64(p.delayMs), metric.WithAttributes(attribute.String("host", host))) - } - return nil - }, fdsOpen, peerDelay); err != nil { - return nil, fmt.Errorf("register glutton observable callback: %w", err) - } - - return s, nil -} - -// Close cancels every running gossip goroutine and waits for them to exit. -func (s *gluttonService) Close() { - s.mu.Lock() - peers := s.peers - s.peers = make(map[string]*peerGossip) - s.mu.Unlock() - for _, p := range peers { - p.cancel() - <-p.done - } -} - -// Write to RAM, either overwriting previously-used RAM or allocating additional RAM -// per request instructions. Data written will be random bytes. -func (s *gluttonService) WriteRAM(ctx context.Context, req *glutton.WriteRAMRequest) (*glutton.WriteRAMResponse, error) { - if req.GetKey() == "" { - return nil, status.Error(codes.InvalidArgument, "key is required") - } - sizeBytes, err := parseBytes(req.GetSize()) + handler, err := glutton.Handler(*mode, svc) if err != nil { - return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) - } - if sizeBytes < 0 { - return nil, status.Error(codes.InvalidArgument, "size must be non-negative") - } - size := int(sizeBytes) - - switch req.GetWriteMode() { - case glutton.WriteMode_WRITE_MODE_TRUNCATE: - buf, err := randomBytes(size) - if err != nil { - return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - s.mu.Lock() - s.ram[req.GetKey()] = buf - delete(s.ramCursor, req.GetKey()) - s.mu.Unlock() - case glutton.WriteMode_WRITE_MODE_OVERWRITE: - s.mu.Lock() - existing := s.ram[req.GetKey()] - if size > len(existing) { - existing = make([]byte, size) - s.ram[req.GetKey()] = existing - delete(s.ramCursor, req.GetKey()) - } - if _, err := rand.Read(existing[:size]); err != nil { - s.mu.Unlock() - return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - s.mu.Unlock() - case glutton.WriteMode_WRITE_MODE_OVERWRITE_ROTATE: - if err := s.rotateRAM(req.GetKey(), size); err != nil { - return nil, err - } - default: - return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) - } - - s.ramWriteBytes.Add(ctx, int64(size)) - return &glutton.WriteRAMResponse{}, nil -} - -// rotateRAM re-randomizes size bytes starting at the key's cursor, wrapping -// at the end of the array, then advances the cursor past the write. Repeated -// rotates therefore walk the whole array instead of re-dirtying the same -// prefix. The cursor lives in process memory, so it rides along in snapshots -// and the walk keeps advancing across suspend/resume cycles. -func (s *gluttonService) rotateRAM(key string, size int) error { - s.mu.Lock() - defer s.mu.Unlock() - existing := s.ram[key] - if len(existing) == 0 { - return status.Errorf(codes.NotFound, "rotate needs an existing array %q; fill with TRUNCATE first", key) - } - if size > len(existing) { - size = len(existing) - } - start := s.ramCursor[key] - head := existing[start:min(start+size, len(existing))] - if _, err := rand.Read(head); err != nil { - return status.Errorf(codes.Internal, "generate random bytes: %v", err) + serverboot.Fatal(ctx, "Invalid --mode", err) } - if wrapped := size - len(head); wrapped > 0 { - if _, err := rand.Read(existing[:wrapped]); err != nil { - return status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - } - s.ramCursor[key] = (start + size) % len(existing) - return nil -} - -// pageSize is the stride of the ReadRAM walk: one byte per 4KiB page is -// enough to force every page resident without the cost of reading them all. -const pageSize = 4096 - -// Walk RAM previously written by WriteRAM, reading one byte per page so -// every touched page must be resident before the response returns. After a -// demand-paged restore this converts restore-time laziness into measurable -// read latency. -func (s *gluttonService) ReadRAM(ctx context.Context, req *glutton.ReadRAMRequest) (*glutton.ReadRAMResponse, error) { - if req.GetKey() == "" { - return nil, status.Error(codes.InvalidArgument, "key is required") - } - s.mu.Lock() - defer s.mu.Unlock() - arr, ok := s.ram[req.GetKey()] - if !ok { - return nil, status.Errorf(codes.NotFound, "no RAM array %q", req.GetKey()) - } - walk := int64(len(arr)) - if req.GetSize() != "" { - n, err := parseBytes(req.GetSize()) - if err != nil { - return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) - } - if n < 0 { - return nil, status.Error(codes.InvalidArgument, "size must be non-negative") - } - walk = min(n, walk) - } - var sum uint32 - for i := int64(0); i < walk; i += pageSize { - sum ^= uint32(arr[i]) - } - s.ramReadBytes.Add(ctx, walk) - return &glutton.ReadRAMResponse{Size: walk, Checksum: sum}, nil -} - -// Write to disk using the specified mode. Data written will be random bytes. -func (s *gluttonService) WriteDisk(ctx context.Context, req *glutton.WriteDiskRequest) (*glutton.WriteDiskResponse, error) { - if !diskKeyRE.MatchString(req.GetKey()) { - return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) - } - if req.GetSize() < 0 { - return nil, status.Error(codes.InvalidArgument, "size must be non-negative") - } - - path := filepath.Join(s.dataDir, req.GetKey()) - - var flag int - switch req.GetWriteMode() { - case glutton.WriteMode_WRITE_MODE_TRUNCATE: - flag = os.O_RDWR | os.O_CREATE | os.O_TRUNC - case glutton.WriteMode_WRITE_MODE_OVERWRITE: - // No O_TRUNC: writes go from offset 0 but any bytes beyond size remain. - flag = os.O_RDWR | os.O_CREATE - default: - return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) - } - - f, err := os.OpenFile(path, flag, 0o600) - if err != nil { - return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) - } - defer f.Close() - - h := sha256.New() - size := int64(req.GetSize()) - if err := streamRandomBytes(io.MultiWriter(f, h), size); err != nil { - return nil, status.Errorf(codes.Internal, "write %s: %v", path, err) - } - - // OVERWRITE has no O_TRUNC, bytes from a larger, earlier write will persist. - // The cursor is already at size, so folding the remainder into the - // same digest completes it without re-reading the prefix. - if req.GetWriteMode() == glutton.WriteMode_WRITE_MODE_OVERWRITE { - tail, err := io.Copy(h, f) - if err != nil { - return nil, status.Errorf(codes.Internal, "hash tail %s: %v", path, err) - } - size += tail - } - - s.diskWriteBytes.Add(ctx, int64(req.GetSize())) - return &glutton.WriteDiskResponse{Size: size, Sha256: h.Sum(nil)}, nil -} - -func (s *gluttonService) ReadDisk(ctx context.Context, req *glutton.ReadDiskRequest) (*glutton.ReadDiskResponse, error) { - if !diskKeyRE.MatchString(req.GetKey()) { - return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) - } - - path := filepath.Join(s.dataDir, req.GetKey()) - - f, err := os.Open(path) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return nil, status.Errorf(codes.NotFound, "file %q not found", req.GetKey()) - } - return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) - } - defer f.Close() - - h := sha256.New() - - if req.GetReadMode() == glutton.ReadMode_READ_MODE_DIGEST_ONLY { - n, err := io.Copy(h, f) - if err != nil { - return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) - } - s.diskReadBytes.Add(ctx, n) - return &glutton.ReadDiskResponse{ - Size: n, - Sha256: h.Sum(nil), - }, nil - } - - data, err := io.ReadAll(io.TeeReader(f, h)) - if err != nil { - return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) - } - - s.diskReadBytes.Add(ctx, int64(len(data))) - return &glutton.ReadDiskResponse{ - Size: int64(len(data)), - Sha256: h.Sum(nil), - Data: data, - }, nil -} - -// Make sure it has the specified number of file descriptors open. It will open or -// close file descriptors to hit the desired count (note this count is in addition to the other -// FDs needed to run the process). -func (s *gluttonService) OpenFD(_ context.Context, req *glutton.OpenFDRequest) (*glutton.OpenFDResponse, error) { - if req.GetCount() < 0 { - return nil, status.Error(codes.InvalidArgument, "count must be non-negative") - } - target := int(req.GetCount()) - - s.mu.Lock() - defer s.mu.Unlock() - - for len(s.fds) > target { - last := len(s.fds) - 1 - if err := s.fds[last].Close(); err != nil { - slog.Warn("Failed to close glutton fd", slog.Any("err", err)) - } - s.fds[last] = nil - s.fds = s.fds[:last] - } - for len(s.fds) < target { - f, err := os.Open(os.DevNull) - if err != nil { - return nil, status.Errorf(codes.Internal, "open %s: %v", os.DevNull, err) - } - s.fds = append(s.fds, f) - } - return &glutton.OpenFDResponse{}, nil -} - -// Receive a ping request, echoing the same response back. -func (s *gluttonService) Ping(ctx context.Context, req *glutton.PingRequest) (*glutton.PingResponse, error) { - s.pingsReceived.Add(ctx, 1) - return &glutton.PingResponse{Message: req.GetMessage()}, nil -} - -// Sends network traffic to a peer glutton. Messages will be sent -// on regular intervals separated by delay_ms. -func (s *gluttonService) Gossip(_ context.Context, req *glutton.GossipRequest) (*glutton.GossipResponse, error) { - want := make(map[string]*glutton.Peer, len(req.GetPeers())) - for _, p := range req.GetPeers() { - if p.GetHost() == "" { - return nil, status.Error(codes.InvalidArgument, "peer host is required") - } - if p.GetDelayMs() <= 0 { - return nil, status.Errorf(codes.InvalidArgument, "peer %q delay_ms must be positive", p.GetHost()) - } - want[p.GetHost()] = p - } - - s.mu.Lock() - var toStop []*peerGossip - for host, existing := range s.peers { - w, ok := want[host] - if !ok || w.GetDelayMs() != existing.delayMs { - toStop = append(toStop, existing) - delete(s.peers, host) - } - } - var toStart []*glutton.Peer - for host, w := range want { - if _, ok := s.peers[host]; !ok { - toStart = append(toStart, w) - } - } - s.mu.Unlock() - - for _, p := range toStop { - p.cancel() - <-p.done - } - - for _, w := range toStart { - gctx, cancel := context.WithCancel(context.Background()) - pg := &peerGossip{ - host: w.GetHost(), - delayMs: w.GetDelayMs(), - cancel: cancel, - done: make(chan struct{}), - } - s.mu.Lock() - s.peers[w.GetHost()] = pg - s.mu.Unlock() - go s.runGossip(gctx, pg) - } - - return &glutton.GossipResponse{}, nil -} - -func (s *gluttonService) runGossip(ctx context.Context, pg *peerGossip) { - defer close(pg.done) - - // grpc.NewClient resolves and connects lazily; the first RPC surfaces - // any failure, so the peer doesn't have to be reachable at start time. - conn, err := grpc.NewClient(pg.host, - grpc.WithTransportCredentials(insecure.NewCredentials()), - grpc.WithStatsHandler(otelgrpc.NewClientHandler()), - ) - if err != nil { - slog.ErrorContext(ctx, "Failed to dial gossip peer", slog.String("host", pg.host), slog.Any("err", err)) - return - } - defer conn.Close() - client := glutton.NewGluttonClient(conn) - - hostAttr := attribute.String("host", pg.host) - ticker := time.NewTicker(time.Duration(pg.delayMs) * time.Millisecond) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - } - msg := uuid.NewString() - start := time.Now() - resp, err := client.Ping(ctx, &glutton.PingRequest{Message: msg}) - latency := time.Since(start).Seconds() - outcome := "ok" - cancelled := err != nil && errors.Is(ctx.Err(), context.Canceled) - switch { - case cancelled: - outcome = "cancelled" - case err != nil: - outcome = "error" - } - attrs := metric.WithAttributes(hostAttr, attribute.String("outcome", outcome)) - s.gossipSent.Add(ctx, 1, attrs) - s.gossipLatency.Record(ctx, latency, attrs) - if cancelled { - return - } - if err != nil { - slog.WarnContext(ctx, "Gossip ping failed", slog.String("host", pg.host), slog.Any("err", err)) - continue - } - if resp.GetMessage() != msg { - slog.WarnContext(ctx, "Gossip ping returned unexpected message", - slog.String("host", pg.host), - slog.String("sent", msg), - slog.String("received", resp.GetMessage()), - ) - } - } -} - -func randomBytes(n int) ([]byte, error) { - buf := make([]byte, n) - if _, err := rand.Read(buf); err != nil { - return nil, err - } - return buf, nil -} - -// streamRandomBytesChunk caps per-syscall random fill and write size so a -// multi-gigabyte WriteDisk doesn't have to materialize in RAM. -const streamRandomBytesChunk = 1 << 20 // 1 MiB - -// streamRandomBytes writes total random bytes to w sequentially, in -// streamRandomBytesChunk-sized chunks. The caller is responsible for the -// file's open mode and starting offset; this writes from the current -// position forward. -func streamRandomBytes(w io.Writer, total int64) error { - if total <= 0 { - return nil - } - buf := make([]byte, streamRandomBytesChunk) - var written int64 - for written < total { - chunk := buf - if remaining := total - written; remaining < int64(len(chunk)) { - chunk = buf[:remaining] - } - if _, err := rand.Read(chunk); err != nil { - return fmt.Errorf("generate random bytes: %w", err) - } - n, err := w.Write(chunk) - if err != nil { - return err - } - written += int64(n) + if err := glutton.NewServer(handler).Serve(lis); err != nil { + serverboot.Fatal(ctx, "Failed to serve", err) } - return nil } diff --git a/internal/benchmarking/glutton/glutton.go b/internal/benchmarking/glutton/glutton.go new file mode 100644 index 0000000000..377842f199 --- /dev/null +++ b/internal/benchmarking/glutton/glutton.go @@ -0,0 +1,726 @@ +// Copyright 2026 Google LLC +// +// 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 glutton implements the benchmarking workload served by +// cmd/benchmarking/glutton: an API for consuming RAM, disk, and file +// descriptors, and for gossiping with other glutton instances, over either +// gRPC or protobuf-over-HTTP. See internal/proto/glutton/glutton.proto. +package glutton + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "sync" + "time" + + "github.com/google/uuid" + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/reflection" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + "github.com/agent-substrate/substrate/internal/ateinterceptors" + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +) + +const meterName = "glutton" + +// Handler builds the request handler for the given wire mode. "grpc" serves +// gRPC alongside the readiness probe on a single listener; "http" serves the +// protobuf-over-HTTP route table. An unknown mode comes back as an error so +// the caller decides how to fail. +func Handler(mode string, svc *Service) (http.Handler, error) { + var handler http.Handler + switch mode { + case "grpc": + srv := grpc.NewServer( + grpc.StatsHandler(otelgrpc.NewServerHandler()), + ) + gluttonpb.RegisterGluttonServer(srv, svc) + reflection.Register(srv) + // The readiness probe is an HTTP GET, so gRPC mode serves it next to + // the gRPC handler on the same listener. + handler = splitGRPC(srv, readyzMux()) + case "http": + // otelhttp at the mux level + per-handler span follows + // docs/dev/best-practices/tracing.md: extract incoming context, + // then name the span after the operation in each handler. + handler = otelhttp.NewHandler(newMux(svc), "/") + default: + return nil, fmt.Errorf("must be grpc or http: %q", mode) + } + return handler, nil +} + +// NewServer enables unencrypted HTTP/2 so gRPC works on the plaintext +// listener, alongside HTTP/1.1 for the readyz probe. +func NewServer(handler http.Handler) *http.Server { + protocols := new(http.Protocols) + protocols.SetHTTP1(true) + protocols.SetUnencryptedHTTP2(true) + return &http.Server{Handler: handler, Protocols: protocols} +} + +// splitGRPC serves gRPC and plain HTTP on one listener: requests with a +// gRPC content-type go to grpcSrv, everything else to rest. All glutton +// RPCs are unary, which is what grpc.Server.ServeHTTP supports. +func splitGRPC(grpcSrv, rest http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") { + grpcSrv.ServeHTTP(w, r) + return + } + rest.ServeHTTP(w, r) + }) +} + +// readyzMux serves the readiness probe both modes need: ateom blocks +// RestoreWorkload until /readyz returns 200, so ResumeActor cannot report +// success before this listener is reachable. +func readyzMux() *http.ServeMux { + mux := http.NewServeMux() + mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + return mux +} + +// newMux builds the HTTP-mode route table on top of the readiness probe. +func newMux(svc *Service) *http.ServeMux { + mux := readyzMux() + mux.HandleFunc("/ping", protoRoute("Ping", svc.Ping)) + mux.HandleFunc("/writedisk", protoRoute("WriteDisk", svc.WriteDisk)) + mux.HandleFunc("/readdisk", protoRoute("ReadDisk", svc.ReadDisk)) + mux.HandleFunc("/writeram", protoRoute("WriteRAM", svc.WriteRAM)) + mux.HandleFunc("/readram", protoRoute("ReadRAM", svc.ReadRAM)) + return mux +} + +// protoRoute wraps a protobuf handler with POST-only routing, protobuf +// unmarshaling, status code mapping, and server-timing headers. +func protoRoute[Req any, Resp proto.Message, PtrReq interface { + *Req + proto.Message +}](spanName string, handler func(context.Context, PtrReq) (Resp, error)) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + start := time.Now() + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + var req Req + ptrReq := PtrReq(&req) + if err := proto.Unmarshal(body, ptrReq); err != nil { + http.Error(w, "unmarshal: "+err.Error(), http.StatusBadRequest) + return + } + ctx, span := otel.Tracer("glutton").Start(r.Context(), spanName) + defer span.End() + resp, err := handler(ctx, ptrReq) + if err != nil { + if st, ok := status.FromError(err); ok { + switch st.Code() { + case codes.InvalidArgument: + http.Error(w, st.Message(), http.StatusBadRequest) + case codes.NotFound: + http.Error(w, st.Message(), http.StatusNotFound) + default: + http.Error(w, st.Message(), http.StatusInternalServerError) + } + } else { + http.Error(w, err.Error(), http.StatusInternalServerError) + } + return + } + out, err := proto.Marshal(resp) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + // Glutton does not run ateinterceptors, so without this the serve path has no + // server-side timing at all. Mirrors the control-plane gRPC trailer so boomer's + // elapsedFromMD logic (source=server) works identically over HTTP. + w.Header().Set(ateinterceptors.ServerElapsedTrailer, + strconv.FormatInt(time.Since(start).Microseconds(), 10)) + w.Header().Set("Content-Type", "application/x-protobuf") + _, _ = w.Write(out) + } +} + +// diskKeyRE rejects anything that could escape the data dir or hit a +// hidden file: only alphanumerics, underscore, and dash are permitted. +var diskKeyRE = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`) + +type Service struct { + gluttonpb.UnimplementedGluttonServer + + dataDir string + + // TODO: split this into per-resource locks (ram, fds, peers). A single + // global mutex serializes unrelated operations across all three. + mu sync.Mutex + ram map[string][]byte + // ramCursor is each array's next WRITE_MODE_OVERWRITE_ROTATE offset. + // Absent means 0; invalidated whenever the array is reallocated. + ramCursor map[string]int + fds []*os.File + peers map[string]*peerGossip + + ramWriteBytes metric.Int64Counter + ramReadBytes metric.Int64Counter + diskWriteBytes metric.Int64Counter + diskReadBytes metric.Int64Counter + pingsReceived metric.Int64Counter + gossipSent metric.Int64Counter + gossipLatency metric.Float64Histogram +} + +type peerGossip struct { + host string + delayMs int32 + cancel context.CancelFunc + done chan struct{} +} + +// New constructs a Service storing WriteDisk files under dir and registers its +// otel instruments. The caller is responsible for creating dir and for calling +// Close to stop any running gossip goroutines. +func New(dir string) (*Service, error) { + s := &Service{ + dataDir: dir, + ram: make(map[string][]byte), + ramCursor: make(map[string]int), + peers: make(map[string]*peerGossip), + } + + m := otel.Meter(meterName) + + var err error + s.ramWriteBytes, err = m.Int64Counter( + "glutton.ram.write.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes written to RAM via WriteRAM over the process lifetime."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.ram.write.bytes counter: %w", err) + } + s.ramReadBytes, err = m.Int64Counter( + "glutton.ram.read.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes walked by ReadRAM over the process lifetime."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.ram.read.bytes counter: %w", err) + } + s.diskWriteBytes, err = m.Int64Counter( + "glutton.disk.write.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes written to disk via WriteDisk over the process lifetime."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.disk.write.bytes counter: %w", err) + } + s.diskReadBytes, err = m.Int64Counter( + "glutton.disk.read.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes read from disk via ReadDisk over the process lifetime."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.disk.read.bytes counter: %w", err) + } + s.pingsReceived, err = m.Int64Counter( + "glutton.ping.requests", + metric.WithDescription("Number of Ping requests received."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.ping.requests counter: %w", err) + } + s.gossipSent, err = m.Int64Counter( + "glutton.gossip.requests.sent", + metric.WithDescription("Number of gossip Ping requests sent per peer."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.gossip.requests.sent counter: %w", err) + } + s.gossipLatency, err = m.Float64Histogram( + "glutton.gossip.latency", + metric.WithUnit("s"), + metric.WithDescription("Latency of gossip Ping requests per peer."), + metric.WithExplicitBucketBoundaries( + 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, + ), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.gossip.latency histogram: %w", err) + } + + fdsOpen, err := m.Int64ObservableGauge( + "glutton.fds.open", + metric.WithDescription("File descriptors currently held open by OpenFD."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.fds.open gauge: %w", err) + } + peerDelay, err := m.Int64ObservableGauge( + "glutton.gossip.delay", + metric.WithUnit("ms"), + metric.WithDescription("Configured gossip delay per peer."), + ) + if err != nil { + return nil, fmt.Errorf("create glutton.gossip.delay gauge: %w", err) + } + + if _, err := m.RegisterCallback(func(_ context.Context, o metric.Observer) error { + s.mu.Lock() + defer s.mu.Unlock() + o.ObserveInt64(fdsOpen, int64(len(s.fds))) + for host, p := range s.peers { + o.ObserveInt64(peerDelay, int64(p.delayMs), metric.WithAttributes(attribute.String("host", host))) + } + return nil + }, fdsOpen, peerDelay); err != nil { + return nil, fmt.Errorf("register glutton observable callback: %w", err) + } + + return s, nil +} + +// Close cancels every running gossip goroutine and waits for them to exit. +func (s *Service) Close() { + s.mu.Lock() + peers := s.peers + s.peers = make(map[string]*peerGossip) + s.mu.Unlock() + for _, p := range peers { + p.cancel() + <-p.done + } +} + +// Write to RAM, either overwriting previously-used RAM or allocating additional RAM +// per request instructions. Data written will be random bytes. +func (s *Service) WriteRAM(ctx context.Context, req *gluttonpb.WriteRAMRequest) (*gluttonpb.WriteRAMResponse, error) { + if req.GetKey() == "" { + return nil, status.Error(codes.InvalidArgument, "key is required") + } + sizeBytes, err := parseBytes(req.GetSize()) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) + } + if sizeBytes < 0 { + return nil, status.Error(codes.InvalidArgument, "size must be non-negative") + } + size := int(sizeBytes) + + switch req.GetWriteMode() { + case gluttonpb.WriteMode_WRITE_MODE_TRUNCATE: + buf, err := randomBytes(size) + if err != nil { + return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + s.mu.Lock() + s.ram[req.GetKey()] = buf + delete(s.ramCursor, req.GetKey()) + s.mu.Unlock() + case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE: + s.mu.Lock() + existing := s.ram[req.GetKey()] + if size > len(existing) { + existing = make([]byte, size) + s.ram[req.GetKey()] = existing + delete(s.ramCursor, req.GetKey()) + } + if _, err := rand.Read(existing[:size]); err != nil { + s.mu.Unlock() + return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + s.mu.Unlock() + case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE_ROTATE: + if err := s.rotateRAM(req.GetKey(), size); err != nil { + return nil, err + } + default: + return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) + } + + s.ramWriteBytes.Add(ctx, int64(size)) + return &gluttonpb.WriteRAMResponse{}, nil +} + +// rotateRAM re-randomizes size bytes starting at the key's cursor, wrapping +// at the end of the array, then advances the cursor past the write. Repeated +// rotates therefore walk the whole array instead of re-dirtying the same +// prefix. The cursor lives in process memory, so it rides along in snapshots +// and the walk keeps advancing across suspend/resume cycles. +func (s *Service) rotateRAM(key string, size int) error { + s.mu.Lock() + defer s.mu.Unlock() + existing := s.ram[key] + if len(existing) == 0 { + return status.Errorf(codes.NotFound, "rotate needs an existing array %q; fill with TRUNCATE first", key) + } + if size > len(existing) { + size = len(existing) + } + start := s.ramCursor[key] + head := existing[start:min(start+size, len(existing))] + if _, err := rand.Read(head); err != nil { + return status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + if wrapped := size - len(head); wrapped > 0 { + if _, err := rand.Read(existing[:wrapped]); err != nil { + return status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + } + s.ramCursor[key] = (start + size) % len(existing) + return nil +} + +// pageSize is the stride of the ReadRAM walk: one byte per 4KiB page is +// enough to force every page resident without the cost of reading them all. +const pageSize = 4096 + +// Walk RAM previously written by WriteRAM, reading one byte per page so +// every touched page must be resident before the response returns. After a +// demand-paged restore this converts restore-time laziness into measurable +// read latency. +func (s *Service) ReadRAM(ctx context.Context, req *gluttonpb.ReadRAMRequest) (*gluttonpb.ReadRAMResponse, error) { + if req.GetKey() == "" { + return nil, status.Error(codes.InvalidArgument, "key is required") + } + s.mu.Lock() + defer s.mu.Unlock() + arr, ok := s.ram[req.GetKey()] + if !ok { + return nil, status.Errorf(codes.NotFound, "no RAM array %q", req.GetKey()) + } + walk := int64(len(arr)) + if req.GetSize() != "" { + n, err := parseBytes(req.GetSize()) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) + } + if n < 0 { + return nil, status.Error(codes.InvalidArgument, "size must be non-negative") + } + walk = min(n, walk) + } + var sum uint32 + for i := int64(0); i < walk; i += pageSize { + sum ^= uint32(arr[i]) + } + s.ramReadBytes.Add(ctx, walk) + return &gluttonpb.ReadRAMResponse{Size: walk, Checksum: sum}, nil +} + +// Write to disk using the specified mode. Data written will be random bytes. +func (s *Service) WriteDisk(ctx context.Context, req *gluttonpb.WriteDiskRequest) (*gluttonpb.WriteDiskResponse, error) { + if !diskKeyRE.MatchString(req.GetKey()) { + return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) + } + if req.GetSize() < 0 { + return nil, status.Error(codes.InvalidArgument, "size must be non-negative") + } + + path := filepath.Join(s.dataDir, req.GetKey()) + + var flag int + switch req.GetWriteMode() { + case gluttonpb.WriteMode_WRITE_MODE_TRUNCATE: + flag = os.O_RDWR | os.O_CREATE | os.O_TRUNC + case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE: + // No O_TRUNC: writes go from offset 0 but any bytes beyond size remain. + flag = os.O_RDWR | os.O_CREATE + default: + return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) + } + + f, err := os.OpenFile(path, flag, 0o600) + if err != nil { + return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) + } + defer f.Close() + + h := sha256.New() + size := int64(req.GetSize()) + if err := streamRandomBytes(io.MultiWriter(f, h), size); err != nil { + return nil, status.Errorf(codes.Internal, "write %s: %v", path, err) + } + + // OVERWRITE has no O_TRUNC, bytes from a larger, earlier write will persist. + // The cursor is already at size, so folding the remainder into the + // same digest completes it without re-reading the prefix. + if req.GetWriteMode() == gluttonpb.WriteMode_WRITE_MODE_OVERWRITE { + tail, err := io.Copy(h, f) + if err != nil { + return nil, status.Errorf(codes.Internal, "hash tail %s: %v", path, err) + } + size += tail + } + + s.diskWriteBytes.Add(ctx, int64(req.GetSize())) + return &gluttonpb.WriteDiskResponse{Size: size, Sha256: h.Sum(nil)}, nil +} + +func (s *Service) ReadDisk(ctx context.Context, req *gluttonpb.ReadDiskRequest) (*gluttonpb.ReadDiskResponse, error) { + if !diskKeyRE.MatchString(req.GetKey()) { + return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) + } + + path := filepath.Join(s.dataDir, req.GetKey()) + + f, err := os.Open(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil, status.Errorf(codes.NotFound, "file %q not found", req.GetKey()) + } + return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) + } + defer f.Close() + + h := sha256.New() + + if req.GetReadMode() == gluttonpb.ReadMode_READ_MODE_DIGEST_ONLY { + n, err := io.Copy(h, f) + if err != nil { + return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) + } + s.diskReadBytes.Add(ctx, n) + return &gluttonpb.ReadDiskResponse{ + Size: n, + Sha256: h.Sum(nil), + }, nil + } + + data, err := io.ReadAll(io.TeeReader(f, h)) + if err != nil { + return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) + } + + s.diskReadBytes.Add(ctx, int64(len(data))) + return &gluttonpb.ReadDiskResponse{ + Size: int64(len(data)), + Sha256: h.Sum(nil), + Data: data, + }, nil +} + +// Make sure it has the specified number of file descriptors open. It will open or +// close file descriptors to hit the desired count (note this count is in addition to the other +// FDs needed to run the process). +func (s *Service) OpenFD(_ context.Context, req *gluttonpb.OpenFDRequest) (*gluttonpb.OpenFDResponse, error) { + if req.GetCount() < 0 { + return nil, status.Error(codes.InvalidArgument, "count must be non-negative") + } + target := int(req.GetCount()) + + s.mu.Lock() + defer s.mu.Unlock() + + for len(s.fds) > target { + last := len(s.fds) - 1 + if err := s.fds[last].Close(); err != nil { + slog.Warn("Failed to close glutton fd", slog.Any("err", err)) + } + s.fds[last] = nil + s.fds = s.fds[:last] + } + for len(s.fds) < target { + f, err := os.Open(os.DevNull) + if err != nil { + return nil, status.Errorf(codes.Internal, "open %s: %v", os.DevNull, err) + } + s.fds = append(s.fds, f) + } + return &gluttonpb.OpenFDResponse{}, nil +} + +// Receive a ping request, echoing the same response back. +func (s *Service) Ping(ctx context.Context, req *gluttonpb.PingRequest) (*gluttonpb.PingResponse, error) { + s.pingsReceived.Add(ctx, 1) + return &gluttonpb.PingResponse{Message: req.GetMessage()}, nil +} + +// Sends network traffic to a peer glutton. Messages will be sent +// on regular intervals separated by delay_ms. +func (s *Service) Gossip(_ context.Context, req *gluttonpb.GossipRequest) (*gluttonpb.GossipResponse, error) { + want := make(map[string]*gluttonpb.Peer, len(req.GetPeers())) + for _, p := range req.GetPeers() { + if p.GetHost() == "" { + return nil, status.Error(codes.InvalidArgument, "peer host is required") + } + if p.GetDelayMs() <= 0 { + return nil, status.Errorf(codes.InvalidArgument, "peer %q delay_ms must be positive", p.GetHost()) + } + want[p.GetHost()] = p + } + + s.mu.Lock() + var toStop []*peerGossip + for host, existing := range s.peers { + w, ok := want[host] + if !ok || w.GetDelayMs() != existing.delayMs { + toStop = append(toStop, existing) + delete(s.peers, host) + } + } + var toStart []*gluttonpb.Peer + for host, w := range want { + if _, ok := s.peers[host]; !ok { + toStart = append(toStart, w) + } + } + s.mu.Unlock() + + for _, p := range toStop { + p.cancel() + <-p.done + } + + for _, w := range toStart { + gctx, cancel := context.WithCancel(context.Background()) + pg := &peerGossip{ + host: w.GetHost(), + delayMs: w.GetDelayMs(), + cancel: cancel, + done: make(chan struct{}), + } + s.mu.Lock() + s.peers[w.GetHost()] = pg + s.mu.Unlock() + go s.runGossip(gctx, pg) + } + + return &gluttonpb.GossipResponse{}, nil +} + +func (s *Service) runGossip(ctx context.Context, pg *peerGossip) { + defer close(pg.done) + + // grpc.NewClient resolves and connects lazily; the first RPC surfaces + // any failure, so the peer doesn't have to be reachable at start time. + conn, err := grpc.NewClient(pg.host, + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithStatsHandler(otelgrpc.NewClientHandler()), + ) + if err != nil { + slog.ErrorContext(ctx, "Failed to dial gossip peer", slog.String("host", pg.host), slog.Any("err", err)) + return + } + defer conn.Close() + client := gluttonpb.NewGluttonClient(conn) + + hostAttr := attribute.String("host", pg.host) + ticker := time.NewTicker(time.Duration(pg.delayMs) * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + msg := uuid.NewString() + start := time.Now() + resp, err := client.Ping(ctx, &gluttonpb.PingRequest{Message: msg}) + latency := time.Since(start).Seconds() + outcome := "ok" + cancelled := err != nil && errors.Is(ctx.Err(), context.Canceled) + switch { + case cancelled: + outcome = "cancelled" + case err != nil: + outcome = "error" + } + attrs := metric.WithAttributes(hostAttr, attribute.String("outcome", outcome)) + s.gossipSent.Add(ctx, 1, attrs) + s.gossipLatency.Record(ctx, latency, attrs) + if cancelled { + return + } + if err != nil { + slog.WarnContext(ctx, "Gossip ping failed", slog.String("host", pg.host), slog.Any("err", err)) + continue + } + if resp.GetMessage() != msg { + slog.WarnContext(ctx, "Gossip ping returned unexpected message", + slog.String("host", pg.host), + slog.String("sent", msg), + slog.String("received", resp.GetMessage()), + ) + } + } +} + +func randomBytes(n int) ([]byte, error) { + buf := make([]byte, n) + if _, err := rand.Read(buf); err != nil { + return nil, err + } + return buf, nil +} + +// streamRandomBytesChunk caps per-syscall random fill and write size so a +// multi-gigabyte WriteDisk doesn't have to materialize in RAM. +const streamRandomBytesChunk = 1 << 20 // 1 MiB + +// streamRandomBytes writes total random bytes to w sequentially, in +// streamRandomBytesChunk-sized chunks. The caller is responsible for the +// file's open mode and starting offset; this writes from the current +// position forward. +func streamRandomBytes(w io.Writer, total int64) error { + if total <= 0 { + return nil + } + buf := make([]byte, streamRandomBytesChunk) + var written int64 + for written < total { + chunk := buf + if remaining := total - written; remaining < int64(len(chunk)) { + chunk = buf[:remaining] + } + if _, err := rand.Read(chunk); err != nil { + return fmt.Errorf("generate random bytes: %w", err) + } + n, err := w.Write(chunk) + if err != nil { + return err + } + written += int64(n) + } + return nil +} diff --git a/cmd/benchmarking/glutton/main_test.go b/internal/benchmarking/glutton/glutton_test.go similarity index 85% rename from cmd/benchmarking/glutton/main_test.go rename to internal/benchmarking/glutton/glutton_test.go index 34d8db2710..8b7bc42934 100644 --- a/cmd/benchmarking/glutton/main_test.go +++ b/internal/benchmarking/glutton/glutton_test.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package glutton import ( "bytes" @@ -28,7 +28,7 @@ import ( "time" "github.com/agent-substrate/substrate/internal/ateinterceptors" - "github.com/agent-substrate/substrate/internal/proto/glutton" + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" @@ -40,14 +40,14 @@ import ( // on a real listener and exercises both protocols against it: the readyz // probe is a plain HTTP GET, and it must not stop gRPC from being served. func TestSplitGRPCServesReadyzAndGRPCOnOneListener(t *testing.T) { - svc, err := newGluttonService(t.TempDir()) + svc, err := New(t.TempDir()) if err != nil { - t.Fatalf("newGluttonService: %v", err) + t.Fatalf("New: %v", err) } defer svc.Close() grpcSrv := grpc.NewServer() - glutton.RegisterGluttonServer(grpcSrv, svc) + gluttonpb.RegisterGluttonServer(grpcSrv, svc) mux := http.NewServeMux() mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { @@ -58,7 +58,7 @@ func TestSplitGRPCServesReadyzAndGRPCOnOneListener(t *testing.T) { if err != nil { t.Fatalf("listen: %v", err) } - srv := newServer(splitGRPC(grpcSrv, mux)) + srv := NewServer(splitGRPC(grpcSrv, mux)) go srv.Serve(lis) defer srv.Close() @@ -81,7 +81,7 @@ func TestSplitGRPCServesReadyzAndGRPCOnOneListener(t *testing.T) { } defer conn.Close() - pong, err := glutton.NewGluttonClient(conn).Ping(ctx, &glutton.PingRequest{Message: "hi"}) + pong, err := gluttonpb.NewGluttonClient(conn).Ping(ctx, &gluttonpb.PingRequest{Message: "hi"}) if err != nil { t.Fatalf("Ping over gRPC: %v", err) } @@ -104,7 +104,7 @@ func TestSplitGRPCRoutesOnContentType(t *testing.T) { if err != nil { t.Fatalf("listen: %v", err) } - srv := newServer(handler) + srv := NewServer(handler) go srv.Serve(lis) defer srv.Close() @@ -123,7 +123,7 @@ func TestSplitGRPCRoutesOnContentType(t *testing.T) { func TestWriteDiskReadDiskRoundTrip(t *testing.T) { tempDir := t.TempDir() - svc, err := newGluttonService(tempDir) + svc, err := New(tempDir) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } @@ -142,10 +142,10 @@ func TestWriteDiskReadDiskRoundTrip(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - writeResp, err := svc.WriteDisk(ctx, &glutton.WriteDiskRequest{ + writeResp, err := svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ Key: tt.key, Size: tt.size, - WriteMode: glutton.WriteMode_WRITE_MODE_TRUNCATE, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, }) if err != nil { t.Fatalf("WriteDisk failed: %v", err) @@ -155,9 +155,9 @@ func TestWriteDiskReadDiskRoundTrip(t *testing.T) { } // 1. Full data read - readResp, err := svc.ReadDisk(ctx, &glutton.ReadDiskRequest{ + readResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ Key: tt.key, - ReadMode: glutton.ReadMode_READ_MODE_DATA, + ReadMode: gluttonpb.ReadMode_READ_MODE_DATA, }) if err != nil { t.Fatalf("ReadDisk (DATA) failed: %v", err) @@ -179,9 +179,9 @@ func TestWriteDiskReadDiskRoundTrip(t *testing.T) { } // 2. Digest-only read - digestResp, err := svc.ReadDisk(ctx, &glutton.ReadDiskRequest{ + digestResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ Key: tt.key, - ReadMode: glutton.ReadMode_READ_MODE_DIGEST_ONLY, + ReadMode: gluttonpb.ReadMode_READ_MODE_DIGEST_ONLY, }) if err != nil { t.Fatalf("ReadDisk (DIGEST_ONLY) failed: %v", err) @@ -201,7 +201,7 @@ func TestWriteDiskReadDiskRoundTrip(t *testing.T) { func TestWriteDiskTruncateProducesExactSize(t *testing.T) { tempDir := t.TempDir() - svc, err := newGluttonService(tempDir) + svc, err := New(tempDir) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } @@ -211,10 +211,10 @@ func TestWriteDiskTruncateProducesExactSize(t *testing.T) { key := "testfile" size := int32(2048) - _, err = svc.WriteDisk(ctx, &glutton.WriteDiskRequest{ + _, err = svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ Key: key, Size: size, - WriteMode: glutton.WriteMode_WRITE_MODE_TRUNCATE, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, }) if err != nil { t.Fatalf("WriteDisk failed: %v", err) @@ -232,7 +232,7 @@ func TestWriteDiskTruncateProducesExactSize(t *testing.T) { func TestWriteDiskOverwriteDigestMatchesReadDisk(t *testing.T) { tempDir := t.TempDir() - svc, err := newGluttonService(tempDir) + svc, err := New(tempDir) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } @@ -242,20 +242,20 @@ func TestWriteDiskOverwriteDigestMatchesReadDisk(t *testing.T) { key := "overwrittenfile" // 1. Initial write of large file (4096 bytes) - _, err = svc.WriteDisk(ctx, &glutton.WriteDiskRequest{ + _, err = svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ Key: key, Size: 4096, - WriteMode: glutton.WriteMode_WRITE_MODE_TRUNCATE, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, }) if err != nil { t.Fatalf("WriteDisk (large) failed: %v", err) } // 2. Overwrite prefix with smaller size (1024 bytes) without truncation - overwriteResp, err := svc.WriteDisk(ctx, &glutton.WriteDiskRequest{ + overwriteResp, err := svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ Key: key, Size: 1024, - WriteMode: glutton.WriteMode_WRITE_MODE_OVERWRITE, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_OVERWRITE, }) if err != nil { t.Fatalf("WriteDisk (overwrite) failed: %v", err) @@ -266,9 +266,9 @@ func TestWriteDiskOverwriteDigestMatchesReadDisk(t *testing.T) { } // 3. ReadDisk reads the entire file (4096 bytes) - readResp, err := svc.ReadDisk(ctx, &glutton.ReadDiskRequest{ + readResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ Key: key, - ReadMode: glutton.ReadMode_READ_MODE_DATA, + ReadMode: gluttonpb.ReadMode_READ_MODE_DATA, }) if err != nil { t.Fatalf("ReadDisk failed: %v", err) @@ -284,14 +284,14 @@ func TestWriteDiskOverwriteDigestMatchesReadDisk(t *testing.T) { func TestReadDiskRejectsInvalidKey(t *testing.T) { tempDir := t.TempDir() - svc, err := newGluttonService(tempDir) + svc, err := New(tempDir) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } defer svc.Close() ctx := context.Background() - _, err = svc.ReadDisk(ctx, &glutton.ReadDiskRequest{Key: "../escape"}) + _, err = svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{Key: "../escape"}) if err == nil { t.Error("expected error for invalid key with path traversal, got nil") } @@ -302,14 +302,14 @@ func TestReadDiskRejectsInvalidKey(t *testing.T) { func TestReadDiskNotFound(t *testing.T) { tempDir := t.TempDir() - svc, err := newGluttonService(tempDir) + svc, err := New(tempDir) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } defer svc.Close() ctx := context.Background() - _, err = svc.ReadDisk(ctx, &glutton.ReadDiskRequest{Key: "nonexistent"}) + _, err = svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{Key: "nonexistent"}) if err == nil { t.Error("expected error for nonexistent file, got nil") } @@ -320,7 +320,7 @@ func TestReadDiskNotFound(t *testing.T) { func TestHTTPRoutes(t *testing.T) { tempDir := t.TempDir() - svc, err := newGluttonService(tempDir) + svc, err := New(tempDir) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } @@ -360,7 +360,7 @@ func TestHTTPRoutes(t *testing.T) { res.Body.Close() // 4. POST /ping -> 200 OK & protobuf Content-Type & ServerElapsedTrailer & echo message - pingReqBytes, _ := proto.Marshal(&glutton.PingRequest{Message: "hello"}) + pingReqBytes, _ := proto.Marshal(&gluttonpb.PingRequest{Message: "hello"}) res, err = http.Post(ts.URL+"/ping", "application/x-protobuf", bytes.NewReader(pingReqBytes)) if err != nil { t.Fatalf("POST /ping failed: %v", err) @@ -376,7 +376,7 @@ func TestHTTPRoutes(t *testing.T) { } body, _ := io.ReadAll(res.Body) res.Body.Close() - var pingResp glutton.PingResponse + var pingResp gluttonpb.PingResponse if err := proto.Unmarshal(body, &pingResp); err != nil { t.Fatalf("unmarshal PingResponse failed: %v", err) } @@ -385,10 +385,10 @@ func TestHTTPRoutes(t *testing.T) { } // 5. POST /writedisk -> 200 OK & protobuf Content-Type - writeReqBytes, _ := proto.Marshal(&glutton.WriteDiskRequest{ + writeReqBytes, _ := proto.Marshal(&gluttonpb.WriteDiskRequest{ Key: "httpfile", Size: 512, - WriteMode: glutton.WriteMode_WRITE_MODE_TRUNCATE, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, }) res, err = http.Post(ts.URL+"/writedisk", "application/x-protobuf", bytes.NewReader(writeReqBytes)) if err != nil { @@ -399,7 +399,7 @@ func TestHTTPRoutes(t *testing.T) { } body, _ = io.ReadAll(res.Body) res.Body.Close() - var writeResp glutton.WriteDiskResponse + var writeResp gluttonpb.WriteDiskResponse if err := proto.Unmarshal(body, &writeResp); err != nil { t.Fatalf("unmarshal WriteDiskResponse failed: %v", err) } @@ -408,7 +408,7 @@ func TestHTTPRoutes(t *testing.T) { } // 6. POST /readdisk -> 200 OK & matching size & digest - readReqBytes, _ := proto.Marshal(&glutton.ReadDiskRequest{Key: "httpfile"}) + readReqBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "httpfile"}) res, err = http.Post(ts.URL+"/readdisk", "application/x-protobuf", bytes.NewReader(readReqBytes)) if err != nil { t.Fatalf("POST /readdisk failed: %v", err) @@ -418,7 +418,7 @@ func TestHTTPRoutes(t *testing.T) { } body, _ = io.ReadAll(res.Body) res.Body.Close() - var readResp glutton.ReadDiskResponse + var readResp gluttonpb.ReadDiskResponse if err := proto.Unmarshal(body, &readResp); err != nil { t.Fatalf("unmarshal ReadDiskResponse failed: %v", err) } @@ -430,7 +430,7 @@ func TestHTTPRoutes(t *testing.T) { } // 7. unknown key -> 404 (NotFound mapping) - missBytes, _ := proto.Marshal(&glutton.ReadDiskRequest{Key: "nosuchfile"}) + missBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "nosuchfile"}) res, err = http.Post(ts.URL+"/readdisk", "application/x-protobuf", bytes.NewReader(missBytes)) if err != nil { t.Fatalf("POST /readdisk miss failed: %v", err) @@ -441,7 +441,7 @@ func TestHTTPRoutes(t *testing.T) { res.Body.Close() // 8. traversal key -> 400 (InvalidArgument mapping) - badBytes, _ := proto.Marshal(&glutton.ReadDiskRequest{Key: "../etc/passwd"}) + badBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "../etc/passwd"}) res, err = http.Post(ts.URL+"/readdisk", "application/x-protobuf", bytes.NewReader(badBytes)) if err != nil { t.Fatalf("POST /readdisk bad key failed: %v", err) diff --git a/cmd/benchmarking/glutton/ram_test.go b/internal/benchmarking/glutton/ram_test.go similarity index 78% rename from cmd/benchmarking/glutton/ram_test.go rename to internal/benchmarking/glutton/ram_test.go index 569a24b500..468f565b3f 100644 --- a/cmd/benchmarking/glutton/ram_test.go +++ b/internal/benchmarking/glutton/ram_test.go @@ -12,21 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package glutton import ( "bytes" "context" "testing" - "github.com/agent-substrate/substrate/internal/proto/glutton" + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" ) -func newRAMTestService(t *testing.T) *gluttonService { +func newRAMTestService(t *testing.T) *Service { t.Helper() - svc, err := newGluttonService(t.TempDir()) + svc, err := New(t.TempDir()) if err != nil { t.Fatalf("failed to create glutton service: %v", err) } @@ -34,20 +34,20 @@ func newRAMTestService(t *testing.T) *gluttonService { return svc } -func fillRAM(t *testing.T, svc *gluttonService, key, size string) { +func fillRAM(t *testing.T, svc *Service, key, size string) { t.Helper() - _, err := svc.WriteRAM(context.Background(), &glutton.WriteRAMRequest{ - Key: key, Size: size, WriteMode: glutton.WriteMode_WRITE_MODE_TRUNCATE, + _, err := svc.WriteRAM(context.Background(), &gluttonpb.WriteRAMRequest{ + Key: key, Size: size, WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, }) if err != nil { t.Fatalf("WriteRAM truncate %s (%s): %v", key, size, err) } } -func rotateRAM(t *testing.T, svc *gluttonService, key, size string) { +func rotateRAM(t *testing.T, svc *Service, key, size string) { t.Helper() - _, err := svc.WriteRAM(context.Background(), &glutton.WriteRAMRequest{ - Key: key, Size: size, WriteMode: glutton.WriteMode_WRITE_MODE_OVERWRITE_ROTATE, + _, err := svc.WriteRAM(context.Background(), &gluttonpb.WriteRAMRequest{ + Key: key, Size: size, WriteMode: gluttonpb.WriteMode_WRITE_MODE_OVERWRITE_ROTATE, }) if err != nil { t.Fatalf("WriteRAM rotate %s (%s): %v", key, size, err) @@ -55,7 +55,7 @@ func rotateRAM(t *testing.T, svc *gluttonService, key, size string) { } // ramCopy snapshots the current bytes of a RAM array for change comparison. -func ramCopy(svc *gluttonService, key string) []byte { +func ramCopy(svc *Service, key string) []byte { svc.mu.Lock() defer svc.mu.Unlock() return append([]byte(nil), svc.ram[key]...) @@ -66,7 +66,7 @@ func TestReadRAMWalksArray(t *testing.T) { ctx := context.Background() fillRAM(t, svc, "m", "64Ki") - whole, err := svc.ReadRAM(ctx, &glutton.ReadRAMRequest{Key: "m"}) + whole, err := svc.ReadRAM(ctx, &gluttonpb.ReadRAMRequest{Key: "m"}) if err != nil { t.Fatalf("ReadRAM whole: %v", err) } @@ -74,7 +74,7 @@ func TestReadRAMWalksArray(t *testing.T) { t.Errorf("whole walk size = %d, want %d", whole.GetSize(), 64<<10) } - again, err := svc.ReadRAM(ctx, &glutton.ReadRAMRequest{Key: "m"}) + again, err := svc.ReadRAM(ctx, &gluttonpb.ReadRAMRequest{Key: "m"}) if err != nil { t.Fatalf("ReadRAM repeat: %v", err) } @@ -82,7 +82,7 @@ func TestReadRAMWalksArray(t *testing.T) { t.Errorf("repeat checksum = %d, want %d (walk must be deterministic)", again.GetChecksum(), whole.GetChecksum()) } - partial, err := svc.ReadRAM(ctx, &glutton.ReadRAMRequest{Key: "m", Size: "4Ki"}) + partial, err := svc.ReadRAM(ctx, &gluttonpb.ReadRAMRequest{Key: "m", Size: "4Ki"}) if err != nil { t.Fatalf("ReadRAM partial: %v", err) } @@ -90,7 +90,7 @@ func TestReadRAMWalksArray(t *testing.T) { t.Errorf("partial walk size = %d, want %d", partial.GetSize(), 4<<10) } - clamped, err := svc.ReadRAM(ctx, &glutton.ReadRAMRequest{Key: "m", Size: "1Gi"}) + clamped, err := svc.ReadRAM(ctx, &gluttonpb.ReadRAMRequest{Key: "m", Size: "1Gi"}) if err != nil { t.Fatalf("ReadRAM oversized: %v", err) } @@ -106,13 +106,13 @@ func TestReadRAMErrors(t *testing.T) { tests := []struct { name string - req *glutton.ReadRAMRequest + req *gluttonpb.ReadRAMRequest code codes.Code }{ - {name: "empty key", req: &glutton.ReadRAMRequest{}, code: codes.InvalidArgument}, - {name: "missing key", req: &glutton.ReadRAMRequest{Key: "nope"}, code: codes.NotFound}, - {name: "bad size", req: &glutton.ReadRAMRequest{Key: "m", Size: "lots"}, code: codes.InvalidArgument}, - {name: "negative size", req: &glutton.ReadRAMRequest{Key: "m", Size: "-1"}, code: codes.InvalidArgument}, + {name: "empty key", req: &gluttonpb.ReadRAMRequest{}, code: codes.InvalidArgument}, + {name: "missing key", req: &gluttonpb.ReadRAMRequest{Key: "nope"}, code: codes.NotFound}, + {name: "bad size", req: &gluttonpb.ReadRAMRequest{Key: "m", Size: "lots"}, code: codes.InvalidArgument}, + {name: "negative size", req: &gluttonpb.ReadRAMRequest{Key: "m", Size: "-1"}, code: codes.InvalidArgument}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { @@ -191,8 +191,8 @@ func TestWriteRAMRotateClampsAndKeepsRotating(t *testing.T) { func TestWriteRAMRotateNeedsExistingArray(t *testing.T) { svc := newRAMTestService(t) - _, err := svc.WriteRAM(context.Background(), &glutton.WriteRAMRequest{ - Key: "nope", Size: "4Ki", WriteMode: glutton.WriteMode_WRITE_MODE_OVERWRITE_ROTATE, + _, err := svc.WriteRAM(context.Background(), &gluttonpb.WriteRAMRequest{ + Key: "nope", Size: "4Ki", WriteMode: gluttonpb.WriteMode_WRITE_MODE_OVERWRITE_ROTATE, }) if status.Code(err) != codes.NotFound { t.Errorf("rotate on missing array = %v, want NotFound", err) diff --git a/cmd/benchmarking/glutton/sizes.go b/internal/benchmarking/glutton/sizes.go similarity index 98% rename from cmd/benchmarking/glutton/sizes.go rename to internal/benchmarking/glutton/sizes.go index c0c839b9f4..d9351698e1 100644 --- a/cmd/benchmarking/glutton/sizes.go +++ b/internal/benchmarking/glutton/sizes.go @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package glutton import ( "fmt" diff --git a/cmd/benchmarking/glutton/sizes_test.go b/internal/benchmarking/glutton/sizes_test.go similarity index 77% rename from cmd/benchmarking/glutton/sizes_test.go rename to internal/benchmarking/glutton/sizes_test.go index 14030a921e..a7d0c1497d 100644 --- a/cmd/benchmarking/glutton/sizes_test.go +++ b/internal/benchmarking/glutton/sizes_test.go @@ -12,13 +12,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -package main +package glutton import ( "context" "testing" - "github.com/agent-substrate/substrate/internal/proto/glutton" + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" ) func TestParseBytes(t *testing.T) { @@ -56,14 +56,14 @@ func TestParseBytes(t *testing.T) { } func TestWriteRAMSuffixedSize(t *testing.T) { - svc, err := newGluttonService(t.TempDir()) + svc, err := New(t.TempDir()) if err != nil { - t.Fatalf("newGluttonService: %v", err) + t.Fatalf("New: %v", err) } defer svc.Close() ctx := context.Background() - if _, err := svc.WriteRAM(ctx, &glutton.WriteRAMRequest{Key: "a", Size: "1Mi"}); err != nil { + if _, err := svc.WriteRAM(ctx, &gluttonpb.WriteRAMRequest{Key: "a", Size: "1Mi"}); err != nil { t.Fatalf("WriteRAM(size=1Mi): %v", err) } svc.mu.Lock() @@ -73,13 +73,13 @@ func TestWriteRAMSuffixedSize(t *testing.T) { t.Errorf("ram[a] = %d bytes, want %d", got, 1<<20) } - if _, err := svc.WriteRAM(ctx, &glutton.WriteRAMRequest{Key: "b", Size: "zap"}); err == nil { + if _, err := svc.WriteRAM(ctx, &gluttonpb.WriteRAMRequest{Key: "b", Size: "zap"}); err == nil { t.Error("WriteRAM(size=zap) succeeded, want error") } - if _, err := svc.WriteRAM(ctx, &glutton.WriteRAMRequest{Key: "c", Size: "-1Gi"}); err == nil { + if _, err := svc.WriteRAM(ctx, &gluttonpb.WriteRAMRequest{Key: "c", Size: "-1Gi"}); err == nil { t.Error("WriteRAM(size=-1Gi) succeeded, want error") } - if _, err := svc.WriteRAM(ctx, &glutton.WriteRAMRequest{Key: "d"}); err == nil { + if _, err := svc.WriteRAM(ctx, &gluttonpb.WriteRAMRequest{Key: "d"}); err == nil { t.Error("WriteRAM(no size) succeeded, want error") } } From a4fa68828a2f6c5c87c64b7644730fca24ff5b10 Mon Sep 17 00:00:00 2001 From: Sairaj Pokale Date: Fri, 28 Aug 2026 17:31:37 -0700 Subject: [PATCH 2/2] benchmarking: split the glutton package and share its route constants glutton.go arrived from the previous commit as a single 640-line file, and the routes it serves were spelled out twice: once in the mux, once in the fake that stands in for it in boomer's tests. Split it by concern into server.go, service.go, metrics.go, and gossip.go, leaving glutton.go with the package doc and the constants describing the workload's contract. The instrument block moves out of New into Service.initMetrics. Add Name, ModeGRPC/ModeHTTP, and one constant per route. The fake re-exports the route constants instead of declaring its own, so a renamed path cannot leave the stand-in answering something the actor does not. Name replaces the four "glutton" literals across the tracer scope, meter scope, and the binary's tracing and metrics init. Instrument names stay written out in full, since hack/verify/metrics.sh greps for them. Boomer keeps its own private route constants: it is the client, and pointing it at the server's implementation package would link the gRPC server and otel instruments into a binary that runs neither. The tests split the same way, and gain coverage for Handler's mode selection and the /writeram route. Signed-off-by: Sairaj Pokale --- cmd/benchmarking/glutton/main.go | 6 +- internal/benchmarking/glutton/fake/server.go | 13 +- internal/benchmarking/glutton/glutton.go | 726 +----------------- internal/benchmarking/glutton/glutton_test.go | 453 ----------- internal/benchmarking/glutton/gossip.go | 151 ++++ internal/benchmarking/glutton/metrics.go | 120 +++ internal/benchmarking/glutton/server.go | 162 ++++ internal/benchmarking/glutton/server_test.go | 303 ++++++++ internal/benchmarking/glutton/service.go | 374 +++++++++ internal/benchmarking/glutton/service_test.go | 226 ++++++ 10 files changed, 1368 insertions(+), 1166 deletions(-) delete mode 100644 internal/benchmarking/glutton/glutton_test.go create mode 100644 internal/benchmarking/glutton/gossip.go create mode 100644 internal/benchmarking/glutton/metrics.go create mode 100644 internal/benchmarking/glutton/server.go create mode 100644 internal/benchmarking/glutton/server_test.go create mode 100644 internal/benchmarking/glutton/service.go create mode 100644 internal/benchmarking/glutton/service_test.go diff --git a/cmd/benchmarking/glutton/main.go b/cmd/benchmarking/glutton/main.go index 126270054d..9c01fb99b6 100644 --- a/cmd/benchmarking/glutton/main.go +++ b/cmd/benchmarking/glutton/main.go @@ -35,7 +35,7 @@ var ( listenAddr = pflag.String("grpc-listen-addr", ":8080", "Address and port the server should listen on (name kept for back-compat; serves whatever --mode picks).") metricsListenAddr = pflag.String("metrics-listen-addr", ":9090", "Address and port the Prometheus metrics server should listen on.") dataDir = pflag.String("data-dir", "", "Directory under which WriteDisk files are stored. Required.") - mode = pflag.String("mode", "grpc", "Wire protocol for the main listener: grpc (default) or http.") + mode = pflag.String("mode", glutton.ModeGRPC, "Wire protocol for the main listener: grpc (default) or http.") showVersion = pflag.Bool("version", false, "Print version and exit.") ) @@ -55,7 +55,7 @@ func main() { serverboot.InitLogger() tp, err := serverboot.InitTracing(ctx, serverboot.TracingOptions{ - ServiceName: "glutton", + ServiceName: glutton.Name, Sampling: serverboot.ResolveTraceSampling(ctx, serverboot.ParentNeverSampling()), }) if err != nil { @@ -63,7 +63,7 @@ func main() { } defer serverboot.ShutdownProvider("TracerProvider", tp.Shutdown) - mp, err := serverboot.InitMetrics(ctx, "glutton") + mp, err := serverboot.InitMetrics(ctx, glutton.Name) if err != nil { serverboot.Fatal(ctx, "Failed to initialize metrics", err) } diff --git a/internal/benchmarking/glutton/fake/server.go b/internal/benchmarking/glutton/fake/server.go index d72f2b1a67..caec2c03e3 100644 --- a/internal/benchmarking/glutton/fake/server.go +++ b/internal/benchmarking/glutton/fake/server.go @@ -25,17 +25,18 @@ import ( "testing" "github.com/agent-substrate/substrate/internal/ateinterceptors" + "github.com/agent-substrate/substrate/internal/benchmarking/glutton" gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" "google.golang.org/protobuf/proto" ) -// Routes the fake serves, mirroring glutton's real HTTP mux. Declared here so -// the fake depends on nothing; collapses onto one source when glutton's core moves. +// Routes the fake serves, re-exported from the real server so the stand-in +// cannot answer a path the actor does not. const ( - WriteDiskRoute = "/writedisk" - ReadDiskRoute = "/readdisk" - WriteRAMRoute = "/writeram" - ReadRAMRoute = "/readram" + WriteDiskRoute = glutton.WriteDiskRoute + ReadDiskRoute = glutton.ReadDiskRoute + WriteRAMRoute = glutton.WriteRAMRoute + ReadRAMRoute = glutton.ReadRAMRoute ) // Server is an httptest-backed stand-in for a glutton actor holding one file. diff --git a/internal/benchmarking/glutton/glutton.go b/internal/benchmarking/glutton/glutton.go index 377842f199..aaa615f3d5 100644 --- a/internal/benchmarking/glutton/glutton.go +++ b/internal/benchmarking/glutton/glutton.go @@ -18,709 +18,27 @@ // gRPC or protobuf-over-HTTP. See internal/proto/glutton/glutton.proto. package glutton -import ( - "context" - "crypto/rand" - "crypto/sha256" - "errors" - "fmt" - "io" - "log/slog" - "net/http" - "os" - "path/filepath" - "regexp" - "strconv" - "strings" - "sync" - "time" - - "github.com/google/uuid" - "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" - "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" - "go.opentelemetry.io/otel" - "go.opentelemetry.io/otel/attribute" - "go.opentelemetry.io/otel/metric" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials/insecure" - "google.golang.org/grpc/reflection" - "google.golang.org/grpc/status" - "google.golang.org/protobuf/proto" - - "github.com/agent-substrate/substrate/internal/ateinterceptors" - gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +// Name is the workload's identity to the telemetry SDK: the tracer and meter +// scope, and the service name the binary reports. Instrument names are written +// out in full at each call site so they match docs/metrics/registry/metrics.yaml +// literally. +const Name = "glutton" + +// Wire modes the main listener can serve, selected by the binary's --mode flag. +const ( + ModeGRPC = "grpc" + ModeHTTP = "http" ) -const meterName = "glutton" - -// Handler builds the request handler for the given wire mode. "grpc" serves -// gRPC alongside the readiness probe on a single listener; "http" serves the -// protobuf-over-HTTP route table. An unknown mode comes back as an error so -// the caller decides how to fail. -func Handler(mode string, svc *Service) (http.Handler, error) { - var handler http.Handler - switch mode { - case "grpc": - srv := grpc.NewServer( - grpc.StatsHandler(otelgrpc.NewServerHandler()), - ) - gluttonpb.RegisterGluttonServer(srv, svc) - reflection.Register(srv) - // The readiness probe is an HTTP GET, so gRPC mode serves it next to - // the gRPC handler on the same listener. - handler = splitGRPC(srv, readyzMux()) - case "http": - // otelhttp at the mux level + per-handler span follows - // docs/dev/best-practices/tracing.md: extract incoming context, - // then name the span after the operation in each handler. - handler = otelhttp.NewHandler(newMux(svc), "/") - default: - return nil, fmt.Errorf("must be grpc or http: %q", mode) - } - return handler, nil -} - -// NewServer enables unencrypted HTTP/2 so gRPC works on the plaintext -// listener, alongside HTTP/1.1 for the readyz probe. -func NewServer(handler http.Handler) *http.Server { - protocols := new(http.Protocols) - protocols.SetHTTP1(true) - protocols.SetUnencryptedHTTP2(true) - return &http.Server{Handler: handler, Protocols: protocols} -} - -// splitGRPC serves gRPC and plain HTTP on one listener: requests with a -// gRPC content-type go to grpcSrv, everything else to rest. All glutton -// RPCs are unary, which is what grpc.Server.ServeHTTP supports. -func splitGRPC(grpcSrv, rest http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") { - grpcSrv.ServeHTTP(w, r) - return - } - rest.ServeHTTP(w, r) - }) -} - -// readyzMux serves the readiness probe both modes need: ateom blocks -// RestoreWorkload until /readyz returns 200, so ResumeActor cannot report -// success before this listener is reachable. -func readyzMux() *http.ServeMux { - mux := http.NewServeMux() - mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - }) - return mux -} - -// newMux builds the HTTP-mode route table on top of the readiness probe. -func newMux(svc *Service) *http.ServeMux { - mux := readyzMux() - mux.HandleFunc("/ping", protoRoute("Ping", svc.Ping)) - mux.HandleFunc("/writedisk", protoRoute("WriteDisk", svc.WriteDisk)) - mux.HandleFunc("/readdisk", protoRoute("ReadDisk", svc.ReadDisk)) - mux.HandleFunc("/writeram", protoRoute("WriteRAM", svc.WriteRAM)) - mux.HandleFunc("/readram", protoRoute("ReadRAM", svc.ReadRAM)) - return mux -} - -// protoRoute wraps a protobuf handler with POST-only routing, protobuf -// unmarshaling, status code mapping, and server-timing headers. -func protoRoute[Req any, Resp proto.Message, PtrReq interface { - *Req - proto.Message -}](spanName string, handler func(context.Context, PtrReq) (Resp, error)) http.HandlerFunc { - return func(w http.ResponseWriter, r *http.Request) { - start := time.Now() - if r.Method != http.MethodPost { - http.Error(w, "method not allowed", http.StatusMethodNotAllowed) - return - } - body, err := io.ReadAll(r.Body) - if err != nil { - http.Error(w, err.Error(), http.StatusBadRequest) - return - } - var req Req - ptrReq := PtrReq(&req) - if err := proto.Unmarshal(body, ptrReq); err != nil { - http.Error(w, "unmarshal: "+err.Error(), http.StatusBadRequest) - return - } - ctx, span := otel.Tracer("glutton").Start(r.Context(), spanName) - defer span.End() - resp, err := handler(ctx, ptrReq) - if err != nil { - if st, ok := status.FromError(err); ok { - switch st.Code() { - case codes.InvalidArgument: - http.Error(w, st.Message(), http.StatusBadRequest) - case codes.NotFound: - http.Error(w, st.Message(), http.StatusNotFound) - default: - http.Error(w, st.Message(), http.StatusInternalServerError) - } - } else { - http.Error(w, err.Error(), http.StatusInternalServerError) - } - return - } - out, err := proto.Marshal(resp) - if err != nil { - http.Error(w, err.Error(), http.StatusInternalServerError) - return - } - // Glutton does not run ateinterceptors, so without this the serve path has no - // server-side timing at all. Mirrors the control-plane gRPC trailer so boomer's - // elapsedFromMD logic (source=server) works identically over HTTP. - w.Header().Set(ateinterceptors.ServerElapsedTrailer, - strconv.FormatInt(time.Since(start).Microseconds(), 10)) - w.Header().Set("Content-Type", "application/x-protobuf") - _, _ = w.Write(out) - } -} - -// diskKeyRE rejects anything that could escape the data dir or hit a -// hidden file: only alphanumerics, underscore, and dash are permitted. -var diskKeyRE = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`) - -type Service struct { - gluttonpb.UnimplementedGluttonServer - - dataDir string - - // TODO: split this into per-resource locks (ram, fds, peers). A single - // global mutex serializes unrelated operations across all three. - mu sync.Mutex - ram map[string][]byte - // ramCursor is each array's next WRITE_MODE_OVERWRITE_ROTATE offset. - // Absent means 0; invalidated whenever the array is reallocated. - ramCursor map[string]int - fds []*os.File - peers map[string]*peerGossip - - ramWriteBytes metric.Int64Counter - ramReadBytes metric.Int64Counter - diskWriteBytes metric.Int64Counter - diskReadBytes metric.Int64Counter - pingsReceived metric.Int64Counter - gossipSent metric.Int64Counter - gossipLatency metric.Float64Histogram -} - -type peerGossip struct { - host string - delayMs int32 - cancel context.CancelFunc - done chan struct{} -} - -// New constructs a Service storing WriteDisk files under dir and registers its -// otel instruments. The caller is responsible for creating dir and for calling -// Close to stop any running gossip goroutines. -func New(dir string) (*Service, error) { - s := &Service{ - dataDir: dir, - ram: make(map[string][]byte), - ramCursor: make(map[string]int), - peers: make(map[string]*peerGossip), - } - - m := otel.Meter(meterName) - - var err error - s.ramWriteBytes, err = m.Int64Counter( - "glutton.ram.write.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes written to RAM via WriteRAM over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.ram.write.bytes counter: %w", err) - } - s.ramReadBytes, err = m.Int64Counter( - "glutton.ram.read.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes walked by ReadRAM over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.ram.read.bytes counter: %w", err) - } - s.diskWriteBytes, err = m.Int64Counter( - "glutton.disk.write.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes written to disk via WriteDisk over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.disk.write.bytes counter: %w", err) - } - s.diskReadBytes, err = m.Int64Counter( - "glutton.disk.read.bytes", - metric.WithUnit("By"), - metric.WithDescription("Total bytes read from disk via ReadDisk over the process lifetime."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.disk.read.bytes counter: %w", err) - } - s.pingsReceived, err = m.Int64Counter( - "glutton.ping.requests", - metric.WithDescription("Number of Ping requests received."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.ping.requests counter: %w", err) - } - s.gossipSent, err = m.Int64Counter( - "glutton.gossip.requests.sent", - metric.WithDescription("Number of gossip Ping requests sent per peer."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.gossip.requests.sent counter: %w", err) - } - s.gossipLatency, err = m.Float64Histogram( - "glutton.gossip.latency", - metric.WithUnit("s"), - metric.WithDescription("Latency of gossip Ping requests per peer."), - metric.WithExplicitBucketBoundaries( - 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, - ), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.gossip.latency histogram: %w", err) - } - - fdsOpen, err := m.Int64ObservableGauge( - "glutton.fds.open", - metric.WithDescription("File descriptors currently held open by OpenFD."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.fds.open gauge: %w", err) - } - peerDelay, err := m.Int64ObservableGauge( - "glutton.gossip.delay", - metric.WithUnit("ms"), - metric.WithDescription("Configured gossip delay per peer."), - ) - if err != nil { - return nil, fmt.Errorf("create glutton.gossip.delay gauge: %w", err) - } - - if _, err := m.RegisterCallback(func(_ context.Context, o metric.Observer) error { - s.mu.Lock() - defer s.mu.Unlock() - o.ObserveInt64(fdsOpen, int64(len(s.fds))) - for host, p := range s.peers { - o.ObserveInt64(peerDelay, int64(p.delayMs), metric.WithAttributes(attribute.String("host", host))) - } - return nil - }, fdsOpen, peerDelay); err != nil { - return nil, fmt.Errorf("register glutton observable callback: %w", err) - } - - return s, nil -} - -// Close cancels every running gossip goroutine and waits for them to exit. -func (s *Service) Close() { - s.mu.Lock() - peers := s.peers - s.peers = make(map[string]*peerGossip) - s.mu.Unlock() - for _, p := range peers { - p.cancel() - <-p.done - } -} - -// Write to RAM, either overwriting previously-used RAM or allocating additional RAM -// per request instructions. Data written will be random bytes. -func (s *Service) WriteRAM(ctx context.Context, req *gluttonpb.WriteRAMRequest) (*gluttonpb.WriteRAMResponse, error) { - if req.GetKey() == "" { - return nil, status.Error(codes.InvalidArgument, "key is required") - } - sizeBytes, err := parseBytes(req.GetSize()) - if err != nil { - return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) - } - if sizeBytes < 0 { - return nil, status.Error(codes.InvalidArgument, "size must be non-negative") - } - size := int(sizeBytes) - - switch req.GetWriteMode() { - case gluttonpb.WriteMode_WRITE_MODE_TRUNCATE: - buf, err := randomBytes(size) - if err != nil { - return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - s.mu.Lock() - s.ram[req.GetKey()] = buf - delete(s.ramCursor, req.GetKey()) - s.mu.Unlock() - case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE: - s.mu.Lock() - existing := s.ram[req.GetKey()] - if size > len(existing) { - existing = make([]byte, size) - s.ram[req.GetKey()] = existing - delete(s.ramCursor, req.GetKey()) - } - if _, err := rand.Read(existing[:size]); err != nil { - s.mu.Unlock() - return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - s.mu.Unlock() - case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE_ROTATE: - if err := s.rotateRAM(req.GetKey(), size); err != nil { - return nil, err - } - default: - return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) - } - - s.ramWriteBytes.Add(ctx, int64(size)) - return &gluttonpb.WriteRAMResponse{}, nil -} - -// rotateRAM re-randomizes size bytes starting at the key's cursor, wrapping -// at the end of the array, then advances the cursor past the write. Repeated -// rotates therefore walk the whole array instead of re-dirtying the same -// prefix. The cursor lives in process memory, so it rides along in snapshots -// and the walk keeps advancing across suspend/resume cycles. -func (s *Service) rotateRAM(key string, size int) error { - s.mu.Lock() - defer s.mu.Unlock() - existing := s.ram[key] - if len(existing) == 0 { - return status.Errorf(codes.NotFound, "rotate needs an existing array %q; fill with TRUNCATE first", key) - } - if size > len(existing) { - size = len(existing) - } - start := s.ramCursor[key] - head := existing[start:min(start+size, len(existing))] - if _, err := rand.Read(head); err != nil { - return status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - if wrapped := size - len(head); wrapped > 0 { - if _, err := rand.Read(existing[:wrapped]); err != nil { - return status.Errorf(codes.Internal, "generate random bytes: %v", err) - } - } - s.ramCursor[key] = (start + size) % len(existing) - return nil -} - -// pageSize is the stride of the ReadRAM walk: one byte per 4KiB page is -// enough to force every page resident without the cost of reading them all. -const pageSize = 4096 - -// Walk RAM previously written by WriteRAM, reading one byte per page so -// every touched page must be resident before the response returns. After a -// demand-paged restore this converts restore-time laziness into measurable -// read latency. -func (s *Service) ReadRAM(ctx context.Context, req *gluttonpb.ReadRAMRequest) (*gluttonpb.ReadRAMResponse, error) { - if req.GetKey() == "" { - return nil, status.Error(codes.InvalidArgument, "key is required") - } - s.mu.Lock() - defer s.mu.Unlock() - arr, ok := s.ram[req.GetKey()] - if !ok { - return nil, status.Errorf(codes.NotFound, "no RAM array %q", req.GetKey()) - } - walk := int64(len(arr)) - if req.GetSize() != "" { - n, err := parseBytes(req.GetSize()) - if err != nil { - return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) - } - if n < 0 { - return nil, status.Error(codes.InvalidArgument, "size must be non-negative") - } - walk = min(n, walk) - } - var sum uint32 - for i := int64(0); i < walk; i += pageSize { - sum ^= uint32(arr[i]) - } - s.ramReadBytes.Add(ctx, walk) - return &gluttonpb.ReadRAMResponse{Size: walk, Checksum: sum}, nil -} - -// Write to disk using the specified mode. Data written will be random bytes. -func (s *Service) WriteDisk(ctx context.Context, req *gluttonpb.WriteDiskRequest) (*gluttonpb.WriteDiskResponse, error) { - if !diskKeyRE.MatchString(req.GetKey()) { - return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) - } - if req.GetSize() < 0 { - return nil, status.Error(codes.InvalidArgument, "size must be non-negative") - } - - path := filepath.Join(s.dataDir, req.GetKey()) - - var flag int - switch req.GetWriteMode() { - case gluttonpb.WriteMode_WRITE_MODE_TRUNCATE: - flag = os.O_RDWR | os.O_CREATE | os.O_TRUNC - case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE: - // No O_TRUNC: writes go from offset 0 but any bytes beyond size remain. - flag = os.O_RDWR | os.O_CREATE - default: - return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) - } - - f, err := os.OpenFile(path, flag, 0o600) - if err != nil { - return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) - } - defer f.Close() - - h := sha256.New() - size := int64(req.GetSize()) - if err := streamRandomBytes(io.MultiWriter(f, h), size); err != nil { - return nil, status.Errorf(codes.Internal, "write %s: %v", path, err) - } - - // OVERWRITE has no O_TRUNC, bytes from a larger, earlier write will persist. - // The cursor is already at size, so folding the remainder into the - // same digest completes it without re-reading the prefix. - if req.GetWriteMode() == gluttonpb.WriteMode_WRITE_MODE_OVERWRITE { - tail, err := io.Copy(h, f) - if err != nil { - return nil, status.Errorf(codes.Internal, "hash tail %s: %v", path, err) - } - size += tail - } - - s.diskWriteBytes.Add(ctx, int64(req.GetSize())) - return &gluttonpb.WriteDiskResponse{Size: size, Sha256: h.Sum(nil)}, nil -} - -func (s *Service) ReadDisk(ctx context.Context, req *gluttonpb.ReadDiskRequest) (*gluttonpb.ReadDiskResponse, error) { - if !diskKeyRE.MatchString(req.GetKey()) { - return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) - } - - path := filepath.Join(s.dataDir, req.GetKey()) - - f, err := os.Open(path) - if err != nil { - if errors.Is(err, os.ErrNotExist) { - return nil, status.Errorf(codes.NotFound, "file %q not found", req.GetKey()) - } - return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) - } - defer f.Close() - - h := sha256.New() - - if req.GetReadMode() == gluttonpb.ReadMode_READ_MODE_DIGEST_ONLY { - n, err := io.Copy(h, f) - if err != nil { - return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) - } - s.diskReadBytes.Add(ctx, n) - return &gluttonpb.ReadDiskResponse{ - Size: n, - Sha256: h.Sum(nil), - }, nil - } - - data, err := io.ReadAll(io.TeeReader(f, h)) - if err != nil { - return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) - } - - s.diskReadBytes.Add(ctx, int64(len(data))) - return &gluttonpb.ReadDiskResponse{ - Size: int64(len(data)), - Sha256: h.Sum(nil), - Data: data, - }, nil -} - -// Make sure it has the specified number of file descriptors open. It will open or -// close file descriptors to hit the desired count (note this count is in addition to the other -// FDs needed to run the process). -func (s *Service) OpenFD(_ context.Context, req *gluttonpb.OpenFDRequest) (*gluttonpb.OpenFDResponse, error) { - if req.GetCount() < 0 { - return nil, status.Error(codes.InvalidArgument, "count must be non-negative") - } - target := int(req.GetCount()) - - s.mu.Lock() - defer s.mu.Unlock() - - for len(s.fds) > target { - last := len(s.fds) - 1 - if err := s.fds[last].Close(); err != nil { - slog.Warn("Failed to close glutton fd", slog.Any("err", err)) - } - s.fds[last] = nil - s.fds = s.fds[:last] - } - for len(s.fds) < target { - f, err := os.Open(os.DevNull) - if err != nil { - return nil, status.Errorf(codes.Internal, "open %s: %v", os.DevNull, err) - } - s.fds = append(s.fds, f) - } - return &gluttonpb.OpenFDResponse{}, nil -} - -// Receive a ping request, echoing the same response back. -func (s *Service) Ping(ctx context.Context, req *gluttonpb.PingRequest) (*gluttonpb.PingResponse, error) { - s.pingsReceived.Add(ctx, 1) - return &gluttonpb.PingResponse{Message: req.GetMessage()}, nil -} - -// Sends network traffic to a peer glutton. Messages will be sent -// on regular intervals separated by delay_ms. -func (s *Service) Gossip(_ context.Context, req *gluttonpb.GossipRequest) (*gluttonpb.GossipResponse, error) { - want := make(map[string]*gluttonpb.Peer, len(req.GetPeers())) - for _, p := range req.GetPeers() { - if p.GetHost() == "" { - return nil, status.Error(codes.InvalidArgument, "peer host is required") - } - if p.GetDelayMs() <= 0 { - return nil, status.Errorf(codes.InvalidArgument, "peer %q delay_ms must be positive", p.GetHost()) - } - want[p.GetHost()] = p - } - - s.mu.Lock() - var toStop []*peerGossip - for host, existing := range s.peers { - w, ok := want[host] - if !ok || w.GetDelayMs() != existing.delayMs { - toStop = append(toStop, existing) - delete(s.peers, host) - } - } - var toStart []*gluttonpb.Peer - for host, w := range want { - if _, ok := s.peers[host]; !ok { - toStart = append(toStart, w) - } - } - s.mu.Unlock() - - for _, p := range toStop { - p.cancel() - <-p.done - } - - for _, w := range toStart { - gctx, cancel := context.WithCancel(context.Background()) - pg := &peerGossip{ - host: w.GetHost(), - delayMs: w.GetDelayMs(), - cancel: cancel, - done: make(chan struct{}), - } - s.mu.Lock() - s.peers[w.GetHost()] = pg - s.mu.Unlock() - go s.runGossip(gctx, pg) - } - - return &gluttonpb.GossipResponse{}, nil -} - -func (s *Service) runGossip(ctx context.Context, pg *peerGossip) { - defer close(pg.done) - - // grpc.NewClient resolves and connects lazily; the first RPC surfaces - // any failure, so the peer doesn't have to be reachable at start time. - conn, err := grpc.NewClient(pg.host, - grpc.WithTransportCredentials(insecure.NewCredentials()), - grpc.WithStatsHandler(otelgrpc.NewClientHandler()), - ) - if err != nil { - slog.ErrorContext(ctx, "Failed to dial gossip peer", slog.String("host", pg.host), slog.Any("err", err)) - return - } - defer conn.Close() - client := gluttonpb.NewGluttonClient(conn) - - hostAttr := attribute.String("host", pg.host) - ticker := time.NewTicker(time.Duration(pg.delayMs) * time.Millisecond) - defer ticker.Stop() - - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - } - msg := uuid.NewString() - start := time.Now() - resp, err := client.Ping(ctx, &gluttonpb.PingRequest{Message: msg}) - latency := time.Since(start).Seconds() - outcome := "ok" - cancelled := err != nil && errors.Is(ctx.Err(), context.Canceled) - switch { - case cancelled: - outcome = "cancelled" - case err != nil: - outcome = "error" - } - attrs := metric.WithAttributes(hostAttr, attribute.String("outcome", outcome)) - s.gossipSent.Add(ctx, 1, attrs) - s.gossipLatency.Record(ctx, latency, attrs) - if cancelled { - return - } - if err != nil { - slog.WarnContext(ctx, "Gossip ping failed", slog.String("host", pg.host), slog.Any("err", err)) - continue - } - if resp.GetMessage() != msg { - slog.WarnContext(ctx, "Gossip ping returned unexpected message", - slog.String("host", pg.host), - slog.String("sent", msg), - slog.String("received", resp.GetMessage()), - ) - } - } -} - -func randomBytes(n int) ([]byte, error) { - buf := make([]byte, n) - if _, err := rand.Read(buf); err != nil { - return nil, err - } - return buf, nil -} - -// streamRandomBytesChunk caps per-syscall random fill and write size so a -// multi-gigabyte WriteDisk doesn't have to materialize in RAM. -const streamRandomBytesChunk = 1 << 20 // 1 MiB - -// streamRandomBytes writes total random bytes to w sequentially, in -// streamRandomBytesChunk-sized chunks. The caller is responsible for the -// file's open mode and starting offset; this writes from the current -// position forward. -func streamRandomBytes(w io.Writer, total int64) error { - if total <= 0 { - return nil - } - buf := make([]byte, streamRandomBytesChunk) - var written int64 - for written < total { - chunk := buf - if remaining := total - written; remaining < int64(len(chunk)) { - chunk = buf[:remaining] - } - if _, err := rand.Read(chunk); err != nil { - return fmt.Errorf("generate random bytes: %w", err) - } - n, err := w.Write(chunk) - if err != nil { - return err - } - written += int64(n) - } - return nil -} +// Routes the HTTP-mode mux serves. ReadyzRoute is served in both modes: ateom +// blocks RestoreWorkload until it answers 200, so ResumeActor cannot report +// success before the listener is reachable. The fake in ./fake aliases these, +// which is what keeps the stand-in and the real mux from drifting apart. +const ( + ReadyzRoute = "/readyz" + PingRoute = "/ping" + WriteDiskRoute = "/writedisk" + ReadDiskRoute = "/readdisk" + WriteRAMRoute = "/writeram" + ReadRAMRoute = "/readram" +) diff --git a/internal/benchmarking/glutton/glutton_test.go b/internal/benchmarking/glutton/glutton_test.go deleted file mode 100644 index 8b7bc42934..0000000000 --- a/internal/benchmarking/glutton/glutton_test.go +++ /dev/null @@ -1,453 +0,0 @@ -// Copyright 2026 Google LLC -// -// 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 glutton - -import ( - "bytes" - "context" - "crypto/sha256" - "io" - "net" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "testing" - "time" - - "github.com/agent-substrate/substrate/internal/ateinterceptors" - gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" - "google.golang.org/grpc" - "google.golang.org/grpc/codes" - "google.golang.org/grpc/credentials/insecure" - "google.golang.org/grpc/status" - "google.golang.org/protobuf/proto" -) - -// TestSplitGRPCServesReadyzAndGRPCOnOneListener starts the grpc-mode handler -// on a real listener and exercises both protocols against it: the readyz -// probe is a plain HTTP GET, and it must not stop gRPC from being served. -func TestSplitGRPCServesReadyzAndGRPCOnOneListener(t *testing.T) { - svc, err := New(t.TempDir()) - if err != nil { - t.Fatalf("New: %v", err) - } - defer svc.Close() - - grpcSrv := grpc.NewServer() - gluttonpb.RegisterGluttonServer(grpcSrv, svc) - - mux := http.NewServeMux() - mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) { - w.WriteHeader(http.StatusOK) - }) - - lis, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("listen: %v", err) - } - srv := NewServer(splitGRPC(grpcSrv, mux)) - go srv.Serve(lis) - defer srv.Close() - - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - resp, err := http.Get("http://" + lis.Addr().String() + "/readyz") - if err != nil { - t.Fatalf("GET /readyz: %v", err) - } - resp.Body.Close() - if resp.StatusCode != http.StatusOK { - t.Errorf("GET /readyz = %d, want %d", resp.StatusCode, http.StatusOK) - } - - conn, err := grpc.NewClient(lis.Addr().String(), - grpc.WithTransportCredentials(insecure.NewCredentials())) - if err != nil { - t.Fatalf("grpc.NewClient: %v", err) - } - defer conn.Close() - - pong, err := gluttonpb.NewGluttonClient(conn).Ping(ctx, &gluttonpb.PingRequest{Message: "hi"}) - if err != nil { - t.Fatalf("Ping over gRPC: %v", err) - } - if pong.GetMessage() != "hi" { - t.Errorf("Ping = %q, want %q", pong.GetMessage(), "hi") - } -} - -// TestSplitGRPCRoutesOnContentType pins the routing rule itself: an HTTP/2 -// request is not enough to reach the gRPC server, the content type is what -// decides. -func TestSplitGRPCRoutesOnContentType(t *testing.T) { - grpcHit := false - handler := splitGRPC( - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { grpcHit = true }), - http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusTeapot) }), - ) - - lis, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("listen: %v", err) - } - srv := NewServer(handler) - go srv.Serve(lis) - defer srv.Close() - - resp, err := http.Get("http://" + lis.Addr().String() + "/anything") - if err != nil { - t.Fatalf("GET: %v", err) - } - resp.Body.Close() - if resp.StatusCode != http.StatusTeapot { - t.Errorf("HTTP/1.1 GET = %d, want %d (the non-gRPC handler)", resp.StatusCode, http.StatusTeapot) - } - if grpcHit { - t.Error("HTTP/1.1 GET reached the gRPC handler") - } -} - -func TestWriteDiskReadDiskRoundTrip(t *testing.T) { - tempDir := t.TempDir() - svc, err := New(tempDir) - if err != nil { - t.Fatalf("failed to create glutton service: %v", err) - } - defer svc.Close() - - ctx := context.Background() - tests := []struct { - name string - key string - size int32 - }{ - {name: "zero size", key: "zero", size: 0}, - {name: "small size", key: "small", size: 1024}, - {name: "chunk unaligned size", key: "unaligned", size: (1 << 20) + 1}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - writeResp, err := svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ - Key: tt.key, - Size: tt.size, - WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, - }) - if err != nil { - t.Fatalf("WriteDisk failed: %v", err) - } - if writeResp.GetSize() != int64(tt.size) { - t.Errorf("WriteDisk size mismatch: got %d, want %d", writeResp.GetSize(), tt.size) - } - - // 1. Full data read - readResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ - Key: tt.key, - ReadMode: gluttonpb.ReadMode_READ_MODE_DATA, - }) - if err != nil { - t.Fatalf("ReadDisk (DATA) failed: %v", err) - } - - if readResp.GetSize() != int64(tt.size) { - t.Errorf("ReadDisk size mismatch: got %d, want %d", readResp.GetSize(), tt.size) - } - if !bytes.Equal(readResp.GetSha256(), writeResp.GetSha256()) { - t.Errorf("sha256 mismatch between WriteDisk and ReadDisk") - } - if len(readResp.GetData()) != int(tt.size) { - t.Errorf("ReadDisk data length mismatch: got %d, want %d", len(readResp.GetData()), tt.size) - } - - computedDigest := sha256.Sum256(readResp.GetData()) - if !bytes.Equal(readResp.GetSha256(), computedDigest[:]) { - t.Errorf("ReadDisk returned sha256 does not match computed digest of returned data") - } - - // 2. Digest-only read - digestResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ - Key: tt.key, - ReadMode: gluttonpb.ReadMode_READ_MODE_DIGEST_ONLY, - }) - if err != nil { - t.Fatalf("ReadDisk (DIGEST_ONLY) failed: %v", err) - } - if digestResp.GetSize() != int64(tt.size) { - t.Errorf("ReadDisk (DIGEST_ONLY) size mismatch: got %d, want %d", digestResp.GetSize(), tt.size) - } - if !bytes.Equal(digestResp.GetSha256(), writeResp.GetSha256()) { - t.Errorf("sha256 mismatch between WriteDisk and ReadDisk (DIGEST_ONLY)") - } - if len(digestResp.GetData()) != 0 { - t.Errorf("ReadDisk (DIGEST_ONLY) should not return data payload, got %d bytes", len(digestResp.GetData())) - } - }) - } -} - -func TestWriteDiskTruncateProducesExactSize(t *testing.T) { - tempDir := t.TempDir() - svc, err := New(tempDir) - if err != nil { - t.Fatalf("failed to create glutton service: %v", err) - } - defer svc.Close() - - ctx := context.Background() - key := "testfile" - size := int32(2048) - - _, err = svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ - Key: key, - Size: size, - WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, - }) - if err != nil { - t.Fatalf("WriteDisk failed: %v", err) - } - - filePath := filepath.Join(tempDir, key) - fi, err := os.Stat(filePath) - if err != nil { - t.Fatalf("os.Stat failed: %v", err) - } - if fi.Size() != int64(size) { - t.Errorf("file size on disk mismatch: got %d, want %d", fi.Size(), size) - } -} - -func TestWriteDiskOverwriteDigestMatchesReadDisk(t *testing.T) { - tempDir := t.TempDir() - svc, err := New(tempDir) - if err != nil { - t.Fatalf("failed to create glutton service: %v", err) - } - defer svc.Close() - - ctx := context.Background() - key := "overwrittenfile" - - // 1. Initial write of large file (4096 bytes) - _, err = svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ - Key: key, - Size: 4096, - WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, - }) - if err != nil { - t.Fatalf("WriteDisk (large) failed: %v", err) - } - - // 2. Overwrite prefix with smaller size (1024 bytes) without truncation - overwriteResp, err := svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ - Key: key, - Size: 1024, - WriteMode: gluttonpb.WriteMode_WRITE_MODE_OVERWRITE, - }) - if err != nil { - t.Fatalf("WriteDisk (overwrite) failed: %v", err) - } - - if overwriteResp.GetSize() != 4096 { - t.Errorf("expected WriteDisk under OVERWRITE to report total file size 4096, got %d", overwriteResp.GetSize()) - } - - // 3. ReadDisk reads the entire file (4096 bytes) - readResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ - Key: key, - ReadMode: gluttonpb.ReadMode_READ_MODE_DATA, - }) - if err != nil { - t.Fatalf("ReadDisk failed: %v", err) - } - - if readResp.GetSize() != 4096 { - t.Errorf("expected ReadDisk size 4096, got %d", readResp.GetSize()) - } - if !bytes.Equal(readResp.GetSha256(), overwriteResp.GetSha256()) { - t.Errorf("expected WriteDisk(OVERWRITE) whole-file digest to match ReadDisk digest") - } -} - -func TestReadDiskRejectsInvalidKey(t *testing.T) { - tempDir := t.TempDir() - svc, err := New(tempDir) - if err != nil { - t.Fatalf("failed to create glutton service: %v", err) - } - defer svc.Close() - - ctx := context.Background() - _, err = svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{Key: "../escape"}) - if err == nil { - t.Error("expected error for invalid key with path traversal, got nil") - } - if s, ok := status.FromError(err); !ok || s.Code() != codes.InvalidArgument { - t.Errorf("expected InvalidArgument code, got %v", err) - } -} - -func TestReadDiskNotFound(t *testing.T) { - tempDir := t.TempDir() - svc, err := New(tempDir) - if err != nil { - t.Fatalf("failed to create glutton service: %v", err) - } - defer svc.Close() - - ctx := context.Background() - _, err = svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{Key: "nonexistent"}) - if err == nil { - t.Error("expected error for nonexistent file, got nil") - } - if s, ok := status.FromError(err); !ok || s.Code() != codes.NotFound { - t.Errorf("expected NotFound code, got %v", err) - } -} - -func TestHTTPRoutes(t *testing.T) { - tempDir := t.TempDir() - svc, err := New(tempDir) - if err != nil { - t.Fatalf("failed to create glutton service: %v", err) - } - defer svc.Close() - - ts := httptest.NewServer(newMux(svc)) - defer ts.Close() - - // 1. /readyz GET -> 200 OK - res, err := http.Get(ts.URL + "/readyz") - if err != nil { - t.Fatalf("GET /readyz failed: %v", err) - } - if res.StatusCode != http.StatusOK { - t.Errorf("GET /readyz status: got %d, want 200", res.StatusCode) - } - res.Body.Close() - - // 2. GET on /ping -> 405 Method Not Allowed - res, err = http.Get(ts.URL + "/ping") - if err != nil { - t.Fatalf("GET /ping failed: %v", err) - } - if res.StatusCode != http.StatusMethodNotAllowed { - t.Errorf("GET /ping status: got %d, want 405", res.StatusCode) - } - res.Body.Close() - - // 3. POST bad body -> 400 Bad Request - res, err = http.Post(ts.URL+"/ping", "application/x-protobuf", bytes.NewReader([]byte("garbage"))) - if err != nil { - t.Fatalf("POST /ping garbage failed: %v", err) - } - if res.StatusCode != http.StatusBadRequest { - t.Errorf("POST /ping garbage status: got %d, want 400", res.StatusCode) - } - res.Body.Close() - - // 4. POST /ping -> 200 OK & protobuf Content-Type & ServerElapsedTrailer & echo message - pingReqBytes, _ := proto.Marshal(&gluttonpb.PingRequest{Message: "hello"}) - res, err = http.Post(ts.URL+"/ping", "application/x-protobuf", bytes.NewReader(pingReqBytes)) - if err != nil { - t.Fatalf("POST /ping failed: %v", err) - } - if res.StatusCode != http.StatusOK { - t.Errorf("POST /ping status: got %d, want 200", res.StatusCode) - } - if ct := res.Header.Get("Content-Type"); ct != "application/x-protobuf" { - t.Errorf("POST /ping Content-Type: got %q, want application/x-protobuf", ct) - } - if elapsed := res.Header.Get(ateinterceptors.ServerElapsedTrailer); elapsed == "" { - t.Errorf("POST /ping missing header %q", ateinterceptors.ServerElapsedTrailer) - } - body, _ := io.ReadAll(res.Body) - res.Body.Close() - var pingResp gluttonpb.PingResponse - if err := proto.Unmarshal(body, &pingResp); err != nil { - t.Fatalf("unmarshal PingResponse failed: %v", err) - } - if pingResp.GetMessage() != "hello" { - t.Errorf("PingResponse message: got %q, want 'hello'", pingResp.GetMessage()) - } - - // 5. POST /writedisk -> 200 OK & protobuf Content-Type - writeReqBytes, _ := proto.Marshal(&gluttonpb.WriteDiskRequest{ - Key: "httpfile", - Size: 512, - WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, - }) - res, err = http.Post(ts.URL+"/writedisk", "application/x-protobuf", bytes.NewReader(writeReqBytes)) - if err != nil { - t.Fatalf("POST /writedisk failed: %v", err) - } - if res.StatusCode != http.StatusOK { - t.Errorf("POST /writedisk status: got %d, want 200", res.StatusCode) - } - body, _ = io.ReadAll(res.Body) - res.Body.Close() - var writeResp gluttonpb.WriteDiskResponse - if err := proto.Unmarshal(body, &writeResp); err != nil { - t.Fatalf("unmarshal WriteDiskResponse failed: %v", err) - } - if writeResp.GetSize() != 512 { - t.Errorf("WriteDiskResponse size: got %d, want 512", writeResp.GetSize()) - } - - // 6. POST /readdisk -> 200 OK & matching size & digest - readReqBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "httpfile"}) - res, err = http.Post(ts.URL+"/readdisk", "application/x-protobuf", bytes.NewReader(readReqBytes)) - if err != nil { - t.Fatalf("POST /readdisk failed: %v", err) - } - if res.StatusCode != http.StatusOK { - t.Errorf("POST /readdisk status: got %d, want 200", res.StatusCode) - } - body, _ = io.ReadAll(res.Body) - res.Body.Close() - var readResp gluttonpb.ReadDiskResponse - if err := proto.Unmarshal(body, &readResp); err != nil { - t.Fatalf("unmarshal ReadDiskResponse failed: %v", err) - } - if readResp.GetSize() != 512 { - t.Errorf("ReadDiskResponse size: got %d, want 512", readResp.GetSize()) - } - if !bytes.Equal(readResp.GetSha256(), writeResp.GetSha256()) { - t.Errorf("sha256 mismatch over HTTP between writedisk and readdisk") - } - - // 7. unknown key -> 404 (NotFound mapping) - missBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "nosuchfile"}) - res, err = http.Post(ts.URL+"/readdisk", "application/x-protobuf", bytes.NewReader(missBytes)) - if err != nil { - t.Fatalf("POST /readdisk miss failed: %v", err) - } - if res.StatusCode != http.StatusNotFound { - t.Errorf("POST /readdisk miss status: got %d, want 404", res.StatusCode) - } - res.Body.Close() - - // 8. traversal key -> 400 (InvalidArgument mapping) - badBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "../etc/passwd"}) - res, err = http.Post(ts.URL+"/readdisk", "application/x-protobuf", bytes.NewReader(badBytes)) - if err != nil { - t.Fatalf("POST /readdisk bad key failed: %v", err) - } - if res.StatusCode != http.StatusBadRequest { - t.Errorf("POST /readdisk bad key status: got %d, want 400", res.StatusCode) - } - res.Body.Close() -} diff --git a/internal/benchmarking/glutton/gossip.go b/internal/benchmarking/glutton/gossip.go new file mode 100644 index 0000000000..6af172e72a --- /dev/null +++ b/internal/benchmarking/glutton/gossip.go @@ -0,0 +1,151 @@ +// Copyright 2026 Google LLC +// +// 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 glutton + +import ( + "context" + "errors" + "log/slog" + "time" + + "github.com/google/uuid" + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" + + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +) + +type peerGossip struct { + host string + delayMs int32 + cancel context.CancelFunc + done chan struct{} +} + +// Sends network traffic to a peer glutton. Messages will be sent +// on regular intervals separated by delay_ms. +func (s *Service) Gossip(_ context.Context, req *gluttonpb.GossipRequest) (*gluttonpb.GossipResponse, error) { + want := make(map[string]*gluttonpb.Peer, len(req.GetPeers())) + for _, p := range req.GetPeers() { + if p.GetHost() == "" { + return nil, status.Error(codes.InvalidArgument, "peer host is required") + } + if p.GetDelayMs() <= 0 { + return nil, status.Errorf(codes.InvalidArgument, "peer %q delay_ms must be positive", p.GetHost()) + } + want[p.GetHost()] = p + } + + s.mu.Lock() + var toStop []*peerGossip + for host, existing := range s.peers { + w, ok := want[host] + if !ok || w.GetDelayMs() != existing.delayMs { + toStop = append(toStop, existing) + delete(s.peers, host) + } + } + var toStart []*gluttonpb.Peer + for host, w := range want { + if _, ok := s.peers[host]; !ok { + toStart = append(toStart, w) + } + } + s.mu.Unlock() + + for _, p := range toStop { + p.cancel() + <-p.done + } + + for _, w := range toStart { + gctx, cancel := context.WithCancel(context.Background()) + pg := &peerGossip{ + host: w.GetHost(), + delayMs: w.GetDelayMs(), + cancel: cancel, + done: make(chan struct{}), + } + s.mu.Lock() + s.peers[w.GetHost()] = pg + s.mu.Unlock() + go s.runGossip(gctx, pg) + } + + return &gluttonpb.GossipResponse{}, nil +} + +func (s *Service) runGossip(ctx context.Context, pg *peerGossip) { + defer close(pg.done) + + // grpc.NewClient resolves and connects lazily; the first RPC surfaces + // any failure, so the peer doesn't have to be reachable at start time. + conn, err := grpc.NewClient(pg.host, + grpc.WithTransportCredentials(insecure.NewCredentials()), + grpc.WithStatsHandler(otelgrpc.NewClientHandler()), + ) + if err != nil { + slog.ErrorContext(ctx, "Failed to dial gossip peer", slog.String("host", pg.host), slog.Any("err", err)) + return + } + defer conn.Close() + client := gluttonpb.NewGluttonClient(conn) + + hostAttr := attribute.String("host", pg.host) + ticker := time.NewTicker(time.Duration(pg.delayMs) * time.Millisecond) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + msg := uuid.NewString() + start := time.Now() + resp, err := client.Ping(ctx, &gluttonpb.PingRequest{Message: msg}) + latency := time.Since(start).Seconds() + outcome := "ok" + cancelled := err != nil && errors.Is(ctx.Err(), context.Canceled) + switch { + case cancelled: + outcome = "cancelled" + case err != nil: + outcome = "error" + } + attrs := metric.WithAttributes(hostAttr, attribute.String("outcome", outcome)) + s.gossipSent.Add(ctx, 1, attrs) + s.gossipLatency.Record(ctx, latency, attrs) + if cancelled { + return + } + if err != nil { + slog.WarnContext(ctx, "Gossip ping failed", slog.String("host", pg.host), slog.Any("err", err)) + continue + } + if resp.GetMessage() != msg { + slog.WarnContext(ctx, "Gossip ping returned unexpected message", + slog.String("host", pg.host), + slog.String("sent", msg), + slog.String("received", resp.GetMessage()), + ) + } + } +} diff --git a/internal/benchmarking/glutton/metrics.go b/internal/benchmarking/glutton/metrics.go new file mode 100644 index 0000000000..3498051e9e --- /dev/null +++ b/internal/benchmarking/glutton/metrics.go @@ -0,0 +1,120 @@ +// Copyright 2026 Google LLC +// +// 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 glutton + +import ( + "context" + "fmt" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +// initMetrics builds every instrument the service reports and stores the +// synchronous ones on s. The two observable gauges stay local: only the +// callback registered here reads them. +func (s *Service) initMetrics() error { + m := otel.Meter(Name) + + var err error + s.ramWriteBytes, err = m.Int64Counter( + "glutton.ram.write.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes written to RAM via WriteRAM over the process lifetime."), + ) + if err != nil { + return fmt.Errorf("create glutton.ram.write.bytes counter: %w", err) + } + s.ramReadBytes, err = m.Int64Counter( + "glutton.ram.read.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes walked by ReadRAM over the process lifetime."), + ) + if err != nil { + return fmt.Errorf("create glutton.ram.read.bytes counter: %w", err) + } + s.diskWriteBytes, err = m.Int64Counter( + "glutton.disk.write.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes written to disk via WriteDisk over the process lifetime."), + ) + if err != nil { + return fmt.Errorf("create glutton.disk.write.bytes counter: %w", err) + } + s.diskReadBytes, err = m.Int64Counter( + "glutton.disk.read.bytes", + metric.WithUnit("By"), + metric.WithDescription("Total bytes read from disk via ReadDisk over the process lifetime."), + ) + if err != nil { + return fmt.Errorf("create glutton.disk.read.bytes counter: %w", err) + } + s.pingsReceived, err = m.Int64Counter( + "glutton.ping.requests", + metric.WithDescription("Number of Ping requests received."), + ) + if err != nil { + return fmt.Errorf("create glutton.ping.requests counter: %w", err) + } + s.gossipSent, err = m.Int64Counter( + "glutton.gossip.requests.sent", + metric.WithDescription("Number of gossip Ping requests sent per peer."), + ) + if err != nil { + return fmt.Errorf("create glutton.gossip.requests.sent counter: %w", err) + } + s.gossipLatency, err = m.Float64Histogram( + "glutton.gossip.latency", + metric.WithUnit("s"), + metric.WithDescription("Latency of gossip Ping requests per peer."), + metric.WithExplicitBucketBoundaries( + 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2.5, 5, 10, + ), + ) + if err != nil { + return fmt.Errorf("create glutton.gossip.latency histogram: %w", err) + } + + fdsOpen, err := m.Int64ObservableGauge( + "glutton.fds.open", + metric.WithDescription("File descriptors currently held open by OpenFD."), + ) + if err != nil { + return fmt.Errorf("create glutton.fds.open gauge: %w", err) + } + peerDelay, err := m.Int64ObservableGauge( + "glutton.gossip.delay", + metric.WithUnit("ms"), + metric.WithDescription("Configured gossip delay per peer."), + ) + if err != nil { + return fmt.Errorf("create glutton.gossip.delay gauge: %w", err) + } + + if _, err := m.RegisterCallback(func(_ context.Context, o metric.Observer) error { + s.mu.Lock() + defer s.mu.Unlock() + o.ObserveInt64(fdsOpen, int64(len(s.fds))) + for host, p := range s.peers { + o.ObserveInt64(peerDelay, int64(p.delayMs), metric.WithAttributes(attribute.String("host", host))) + } + return nil + }, fdsOpen, peerDelay); err != nil { + return fmt.Errorf("register glutton observable callback: %w", err) + } + + return nil +} diff --git a/internal/benchmarking/glutton/server.go b/internal/benchmarking/glutton/server.go new file mode 100644 index 0000000000..950cc9be7a --- /dev/null +++ b/internal/benchmarking/glutton/server.go @@ -0,0 +1,162 @@ +// Copyright 2026 Google LLC +// +// 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 glutton + +import ( + "context" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" + + "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" + "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" + "go.opentelemetry.io/otel" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/reflection" + "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" + + "github.com/agent-substrate/substrate/internal/ateinterceptors" + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +) + +// Handler builds the request handler for the given wire mode. ModeGRPC serves +// gRPC alongside the readiness probe on a single listener; ModeHTTP serves the +// protobuf-over-HTTP route table. An unknown mode comes back as an error so +// the caller decides how to fail. +func Handler(mode string, svc *Service) (http.Handler, error) { + var handler http.Handler + switch mode { + case ModeGRPC: + srv := grpc.NewServer( + grpc.StatsHandler(otelgrpc.NewServerHandler()), + ) + gluttonpb.RegisterGluttonServer(srv, svc) + reflection.Register(srv) + // The readiness probe is an HTTP GET, so gRPC mode serves it next to + // the gRPC handler on the same listener. + handler = splitGRPC(srv, readyzMux()) + case ModeHTTP: + // otelhttp at the mux level + per-handler span follows + // docs/dev/best-practices/tracing.md: extract incoming context, + // then name the span after the operation in each handler. + handler = otelhttp.NewHandler(newMux(svc), "/") + default: + return nil, fmt.Errorf("must be %s or %s: %q", ModeGRPC, ModeHTTP, mode) + } + return handler, nil +} + +// NewServer enables unencrypted HTTP/2 so gRPC works on the plaintext +// listener, alongside HTTP/1.1 for the readyz probe. +func NewServer(handler http.Handler) *http.Server { + protocols := new(http.Protocols) + protocols.SetHTTP1(true) + protocols.SetUnencryptedHTTP2(true) + return &http.Server{Handler: handler, Protocols: protocols} +} + +// splitGRPC serves gRPC and plain HTTP on one listener: requests with a +// gRPC content-type go to grpcSrv, everything else to rest. All glutton +// RPCs are unary, which is what grpc.Server.ServeHTTP supports. +func splitGRPC(grpcSrv, rest http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.ProtoMajor == 2 && strings.HasPrefix(r.Header.Get("Content-Type"), "application/grpc") { + grpcSrv.ServeHTTP(w, r) + return + } + rest.ServeHTTP(w, r) + }) +} + +// readyzMux serves the readiness probe both modes need. +func readyzMux() *http.ServeMux { + mux := http.NewServeMux() + mux.HandleFunc(ReadyzRoute, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + return mux +} + +// newMux builds the HTTP-mode route table on top of the readiness probe. +func newMux(svc *Service) *http.ServeMux { + mux := readyzMux() + mux.HandleFunc(PingRoute, protoRoute("Ping", svc.Ping)) + mux.HandleFunc(WriteDiskRoute, protoRoute("WriteDisk", svc.WriteDisk)) + mux.HandleFunc(ReadDiskRoute, protoRoute("ReadDisk", svc.ReadDisk)) + mux.HandleFunc(WriteRAMRoute, protoRoute("WriteRAM", svc.WriteRAM)) + mux.HandleFunc(ReadRAMRoute, protoRoute("ReadRAM", svc.ReadRAM)) + return mux +} + +// protoRoute wraps a protobuf handler with POST-only routing, protobuf +// unmarshaling, status code mapping, and server-timing headers. +func protoRoute[Req any, Resp proto.Message, PtrReq interface { + *Req + proto.Message +}](spanName string, handler func(context.Context, PtrReq) (Resp, error)) http.HandlerFunc { + return func(w http.ResponseWriter, r *http.Request) { + start := time.Now() + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + var req Req + ptrReq := PtrReq(&req) + if err := proto.Unmarshal(body, ptrReq); err != nil { + http.Error(w, "unmarshal: "+err.Error(), http.StatusBadRequest) + return + } + ctx, span := otel.Tracer(Name).Start(r.Context(), spanName) + defer span.End() + resp, err := handler(ctx, ptrReq) + if err != nil { + if st, ok := status.FromError(err); ok { + switch st.Code() { + case codes.InvalidArgument: + http.Error(w, st.Message(), http.StatusBadRequest) + case codes.NotFound: + http.Error(w, st.Message(), http.StatusNotFound) + default: + http.Error(w, st.Message(), http.StatusInternalServerError) + } + } else { + http.Error(w, err.Error(), http.StatusInternalServerError) + } + return + } + out, err := proto.Marshal(resp) + if err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + // Glutton does not run ateinterceptors, so without this the serve path has no + // server-side timing at all. Mirrors the control-plane gRPC trailer so boomer's + // elapsedFromMD logic (source=server) works identically over HTTP. + w.Header().Set(ateinterceptors.ServerElapsedTrailer, + strconv.FormatInt(time.Since(start).Microseconds(), 10)) + w.Header().Set("Content-Type", "application/x-protobuf") + _, _ = w.Write(out) + } +} diff --git a/internal/benchmarking/glutton/server_test.go b/internal/benchmarking/glutton/server_test.go new file mode 100644 index 0000000000..4b04246822 --- /dev/null +++ b/internal/benchmarking/glutton/server_test.go @@ -0,0 +1,303 @@ +// Copyright 2026 Google LLC +// +// 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 glutton + +import ( + "bytes" + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "testing" + "time" + + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" + "google.golang.org/protobuf/proto" + + "github.com/agent-substrate/substrate/internal/ateinterceptors" + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +) + +// TestSplitGRPCServesReadyzAndGRPCOnOneListener starts the grpc-mode handler +// on a real listener and exercises both protocols against it: the readyz +// probe is a plain HTTP GET, and it must not stop gRPC from being served. +func TestSplitGRPCServesReadyzAndGRPCOnOneListener(t *testing.T) { + svc, err := New(t.TempDir()) + if err != nil { + t.Fatalf("New: %v", err) + } + defer svc.Close() + + grpcSrv := grpc.NewServer() + gluttonpb.RegisterGluttonServer(grpcSrv, svc) + + mux := http.NewServeMux() + mux.HandleFunc(ReadyzRoute, func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + }) + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + srv := NewServer(splitGRPC(grpcSrv, mux)) + go srv.Serve(lis) + defer srv.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + resp, err := http.Get("http://" + lis.Addr().String() + ReadyzRoute) + if err != nil { + t.Fatalf("GET %s: %v", ReadyzRoute, err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Errorf("GET %s = %d, want %d", ReadyzRoute, resp.StatusCode, http.StatusOK) + } + + conn, err := grpc.NewClient(lis.Addr().String(), + grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + t.Fatalf("grpc.NewClient: %v", err) + } + defer conn.Close() + + pong, err := gluttonpb.NewGluttonClient(conn).Ping(ctx, &gluttonpb.PingRequest{Message: "hi"}) + if err != nil { + t.Fatalf("Ping over gRPC: %v", err) + } + if pong.GetMessage() != "hi" { + t.Errorf("Ping = %q, want %q", pong.GetMessage(), "hi") + } +} + +// TestSplitGRPCRoutesOnContentType pins the routing rule itself: an HTTP/2 +// request is not enough to reach the gRPC server, the content type is what +// decides. +func TestSplitGRPCRoutesOnContentType(t *testing.T) { + grpcHit := false + handler := splitGRPC( + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { grpcHit = true }), + http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusTeapot) }), + ) + + lis, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + srv := NewServer(handler) + go srv.Serve(lis) + defer srv.Close() + + resp, err := http.Get("http://" + lis.Addr().String() + "/anything") + if err != nil { + t.Fatalf("GET: %v", err) + } + resp.Body.Close() + if resp.StatusCode != http.StatusTeapot { + t.Errorf("HTTP/1.1 GET = %d, want %d (the non-gRPC handler)", resp.StatusCode, http.StatusTeapot) + } + if grpcHit { + t.Error("HTTP/1.1 GET reached the gRPC handler") + } +} + +func TestHandlerRejectsUnknownMode(t *testing.T) { + svc, err := New(t.TempDir()) + if err != nil { + t.Fatalf("New: %v", err) + } + defer svc.Close() + + for _, mode := range []string{ModeGRPC, ModeHTTP} { + if _, err := Handler(mode, svc); err != nil { + t.Errorf("Handler(%q): %v", mode, err) + } + } + if _, err := Handler("quic", svc); err == nil { + t.Error("Handler(\"quic\") succeeded, want error") + } +} + +func TestHTTPRoutes(t *testing.T) { + tempDir := t.TempDir() + svc, err := New(tempDir) + if err != nil { + t.Fatalf("failed to create glutton service: %v", err) + } + defer svc.Close() + + ts := httptest.NewServer(newMux(svc)) + defer ts.Close() + + // 1. /readyz GET -> 200 OK + res, err := http.Get(ts.URL + ReadyzRoute) + if err != nil { + t.Fatalf("GET %s failed: %v", ReadyzRoute, err) + } + if res.StatusCode != http.StatusOK { + t.Errorf("GET %s status: got %d, want 200", ReadyzRoute, res.StatusCode) + } + res.Body.Close() + + // 2. GET on /ping -> 405 Method Not Allowed + res, err = http.Get(ts.URL + PingRoute) + if err != nil { + t.Fatalf("GET %s failed: %v", PingRoute, err) + } + if res.StatusCode != http.StatusMethodNotAllowed { + t.Errorf("GET %s status: got %d, want 405", PingRoute, res.StatusCode) + } + res.Body.Close() + + // 3. POST bad body -> 400 Bad Request + res, err = http.Post(ts.URL+PingRoute, "application/x-protobuf", bytes.NewReader([]byte("garbage"))) + if err != nil { + t.Fatalf("POST %s garbage failed: %v", PingRoute, err) + } + if res.StatusCode != http.StatusBadRequest { + t.Errorf("POST %s garbage status: got %d, want 400", PingRoute, res.StatusCode) + } + res.Body.Close() + + // 4. POST /ping -> 200 OK & protobuf Content-Type & ServerElapsedTrailer & echo message + pingReqBytes, _ := proto.Marshal(&gluttonpb.PingRequest{Message: "hello"}) + res, err = http.Post(ts.URL+PingRoute, "application/x-protobuf", bytes.NewReader(pingReqBytes)) + if err != nil { + t.Fatalf("POST %s failed: %v", PingRoute, err) + } + if res.StatusCode != http.StatusOK { + t.Errorf("POST %s status: got %d, want 200", PingRoute, res.StatusCode) + } + if ct := res.Header.Get("Content-Type"); ct != "application/x-protobuf" { + t.Errorf("POST %s Content-Type: got %q, want application/x-protobuf", PingRoute, ct) + } + if elapsed := res.Header.Get(ateinterceptors.ServerElapsedTrailer); elapsed == "" { + t.Errorf("POST %s missing header %q", PingRoute, ateinterceptors.ServerElapsedTrailer) + } + body, _ := io.ReadAll(res.Body) + res.Body.Close() + var pingResp gluttonpb.PingResponse + if err := proto.Unmarshal(body, &pingResp); err != nil { + t.Fatalf("unmarshal PingResponse failed: %v", err) + } + if pingResp.GetMessage() != "hello" { + t.Errorf("PingResponse message: got %q, want 'hello'", pingResp.GetMessage()) + } + + // 5. POST /writedisk -> 200 OK & protobuf Content-Type + writeReqBytes, _ := proto.Marshal(&gluttonpb.WriteDiskRequest{ + Key: "httpfile", + Size: 512, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, + }) + res, err = http.Post(ts.URL+WriteDiskRoute, "application/x-protobuf", bytes.NewReader(writeReqBytes)) + if err != nil { + t.Fatalf("POST %s failed: %v", WriteDiskRoute, err) + } + if res.StatusCode != http.StatusOK { + t.Errorf("POST %s status: got %d, want 200", WriteDiskRoute, res.StatusCode) + } + body, _ = io.ReadAll(res.Body) + res.Body.Close() + var writeResp gluttonpb.WriteDiskResponse + if err := proto.Unmarshal(body, &writeResp); err != nil { + t.Fatalf("unmarshal WriteDiskResponse failed: %v", err) + } + if writeResp.GetSize() != 512 { + t.Errorf("WriteDiskResponse size: got %d, want 512", writeResp.GetSize()) + } + + // 6. POST /readdisk -> 200 OK & matching size & digest + readReqBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "httpfile"}) + res, err = http.Post(ts.URL+ReadDiskRoute, "application/x-protobuf", bytes.NewReader(readReqBytes)) + if err != nil { + t.Fatalf("POST %s failed: %v", ReadDiskRoute, err) + } + if res.StatusCode != http.StatusOK { + t.Errorf("POST %s status: got %d, want 200", ReadDiskRoute, res.StatusCode) + } + body, _ = io.ReadAll(res.Body) + res.Body.Close() + var readResp gluttonpb.ReadDiskResponse + if err := proto.Unmarshal(body, &readResp); err != nil { + t.Fatalf("unmarshal ReadDiskResponse failed: %v", err) + } + if readResp.GetSize() != 512 { + t.Errorf("ReadDiskResponse size: got %d, want 512", readResp.GetSize()) + } + if !bytes.Equal(readResp.GetSha256(), writeResp.GetSha256()) { + t.Errorf("sha256 mismatch over HTTP between writedisk and readdisk") + } + + // 7. POST /writeram -> 200 OK, reachable in HTTP mode + ramReqBytes, _ := proto.Marshal(&gluttonpb.WriteRAMRequest{ + Key: "httpram", + Size: "1Ki", + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, + }) + res, err = http.Post(ts.URL+WriteRAMRoute, "application/x-protobuf", bytes.NewReader(ramReqBytes)) + if err != nil { + t.Fatalf("POST %s failed: %v", WriteRAMRoute, err) + } + if res.StatusCode != http.StatusOK { + t.Errorf("POST %s status: got %d, want 200", WriteRAMRoute, res.StatusCode) + } + res.Body.Close() + + // 8. POST /readram -> 200 OK & matching size + ramReadReqBytes, _ := proto.Marshal(&gluttonpb.ReadRAMRequest{Key: "httpram"}) + res, err = http.Post(ts.URL+ReadRAMRoute, "application/x-protobuf", bytes.NewReader(ramReadReqBytes)) + if err != nil { + t.Fatalf("POST %s failed: %v", ReadRAMRoute, err) + } + if res.StatusCode != http.StatusOK { + t.Errorf("POST %s status: got %d, want 200", ReadRAMRoute, res.StatusCode) + } + body, _ = io.ReadAll(res.Body) + res.Body.Close() + var ramReadResp gluttonpb.ReadRAMResponse + if err := proto.Unmarshal(body, &ramReadResp); err != nil { + t.Fatalf("unmarshal ReadRAMResponse failed: %v", err) + } + if ramReadResp.GetSize() != 1024 { + t.Errorf("ReadRAMResponse size: got %d, want 1024", ramReadResp.GetSize()) + } + + // 9. unknown key -> 404 (NotFound mapping) + missBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "nosuchfile"}) + res, err = http.Post(ts.URL+ReadDiskRoute, "application/x-protobuf", bytes.NewReader(missBytes)) + if err != nil { + t.Fatalf("POST %s miss failed: %v", ReadDiskRoute, err) + } + if res.StatusCode != http.StatusNotFound { + t.Errorf("POST %s miss status: got %d, want 404", ReadDiskRoute, res.StatusCode) + } + res.Body.Close() + + // 10. traversal key -> 400 (InvalidArgument mapping) + badBytes, _ := proto.Marshal(&gluttonpb.ReadDiskRequest{Key: "../etc/passwd"}) + res, err = http.Post(ts.URL+ReadDiskRoute, "application/x-protobuf", bytes.NewReader(badBytes)) + if err != nil { + t.Fatalf("POST %s bad key failed: %v", ReadDiskRoute, err) + } + if res.StatusCode != http.StatusBadRequest { + t.Errorf("POST %s bad key status: got %d, want 400", ReadDiskRoute, res.StatusCode) + } + res.Body.Close() +} diff --git a/internal/benchmarking/glutton/service.go b/internal/benchmarking/glutton/service.go new file mode 100644 index 0000000000..a66d77843d --- /dev/null +++ b/internal/benchmarking/glutton/service.go @@ -0,0 +1,374 @@ +// Copyright 2026 Google LLC +// +// 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 glutton + +import ( + "context" + "crypto/rand" + "crypto/sha256" + "errors" + "fmt" + "io" + "log/slog" + "os" + "path/filepath" + "regexp" + "sync" + + "go.opentelemetry.io/otel/metric" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +) + +// diskKeyRE rejects anything that could escape the data dir or hit a +// hidden file: only alphanumerics, underscore, and dash are permitted. +var diskKeyRE = regexp.MustCompile(`^[a-zA-Z0-9_-]+$`) + +type Service struct { + gluttonpb.UnimplementedGluttonServer + + dataDir string + + // TODO: split this into per-resource locks (ram, fds, peers). A single + // global mutex serializes unrelated operations across all three. + mu sync.Mutex + ram map[string][]byte + // ramCursor is each array's next WRITE_MODE_OVERWRITE_ROTATE offset. + // Absent means 0; invalidated whenever the array is reallocated. + ramCursor map[string]int + fds []*os.File + peers map[string]*peerGossip + + ramWriteBytes metric.Int64Counter + ramReadBytes metric.Int64Counter + diskWriteBytes metric.Int64Counter + diskReadBytes metric.Int64Counter + pingsReceived metric.Int64Counter + gossipSent metric.Int64Counter + gossipLatency metric.Float64Histogram +} + +// New constructs a Service storing WriteDisk files under dir and registers its +// otel instruments. The caller is responsible for creating dir and for calling +// Close to stop any running gossip goroutines. +func New(dir string) (*Service, error) { + s := &Service{ + dataDir: dir, + ram: make(map[string][]byte), + ramCursor: make(map[string]int), + peers: make(map[string]*peerGossip), + } + if err := s.initMetrics(); err != nil { + return nil, err + } + return s, nil +} + +// Close cancels every running gossip goroutine and waits for them to exit. +func (s *Service) Close() { + s.mu.Lock() + peers := s.peers + s.peers = make(map[string]*peerGossip) + s.mu.Unlock() + for _, p := range peers { + p.cancel() + <-p.done + } +} + +// Write to RAM, either overwriting previously-used RAM or allocating additional RAM +// per request instructions. Data written will be random bytes. +func (s *Service) WriteRAM(ctx context.Context, req *gluttonpb.WriteRAMRequest) (*gluttonpb.WriteRAMResponse, error) { + if req.GetKey() == "" { + return nil, status.Error(codes.InvalidArgument, "key is required") + } + sizeBytes, err := parseBytes(req.GetSize()) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) + } + if sizeBytes < 0 { + return nil, status.Error(codes.InvalidArgument, "size must be non-negative") + } + size := int(sizeBytes) + + switch req.GetWriteMode() { + case gluttonpb.WriteMode_WRITE_MODE_TRUNCATE: + buf, err := randomBytes(size) + if err != nil { + return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + s.mu.Lock() + s.ram[req.GetKey()] = buf + delete(s.ramCursor, req.GetKey()) + s.mu.Unlock() + case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE: + s.mu.Lock() + existing := s.ram[req.GetKey()] + if size > len(existing) { + existing = make([]byte, size) + s.ram[req.GetKey()] = existing + delete(s.ramCursor, req.GetKey()) + } + if _, err := rand.Read(existing[:size]); err != nil { + s.mu.Unlock() + return nil, status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + s.mu.Unlock() + case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE_ROTATE: + if err := s.rotateRAM(req.GetKey(), size); err != nil { + return nil, err + } + default: + return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) + } + + s.ramWriteBytes.Add(ctx, int64(size)) + return &gluttonpb.WriteRAMResponse{}, nil +} + +// rotateRAM re-randomizes size bytes starting at the key's cursor, wrapping +// at the end of the array, then advances the cursor past the write. Repeated +// rotates therefore walk the whole array instead of re-dirtying the same +// prefix. The cursor lives in process memory, so it rides along in snapshots +// and the walk keeps advancing across suspend/resume cycles. +func (s *Service) rotateRAM(key string, size int) error { + s.mu.Lock() + defer s.mu.Unlock() + existing := s.ram[key] + if len(existing) == 0 { + return status.Errorf(codes.NotFound, "rotate needs an existing array %q; fill with TRUNCATE first", key) + } + if size > len(existing) { + size = len(existing) + } + start := s.ramCursor[key] + head := existing[start:min(start+size, len(existing))] + if _, err := rand.Read(head); err != nil { + return status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + if wrapped := size - len(head); wrapped > 0 { + if _, err := rand.Read(existing[:wrapped]); err != nil { + return status.Errorf(codes.Internal, "generate random bytes: %v", err) + } + } + s.ramCursor[key] = (start + size) % len(existing) + return nil +} + +// pageSize is the stride of the ReadRAM walk: one byte per 4KiB page is +// enough to force every page resident without the cost of reading them all. +const pageSize = 4096 + +// Walk RAM previously written by WriteRAM, reading one byte per page so +// every touched page must be resident before the response returns. After a +// demand-paged restore this converts restore-time laziness into measurable +// read latency. +func (s *Service) ReadRAM(ctx context.Context, req *gluttonpb.ReadRAMRequest) (*gluttonpb.ReadRAMResponse, error) { + if req.GetKey() == "" { + return nil, status.Error(codes.InvalidArgument, "key is required") + } + s.mu.Lock() + defer s.mu.Unlock() + arr, ok := s.ram[req.GetKey()] + if !ok { + return nil, status.Errorf(codes.NotFound, "no RAM array %q", req.GetKey()) + } + walk := int64(len(arr)) + if req.GetSize() != "" { + n, err := parseBytes(req.GetSize()) + if err != nil { + return nil, status.Errorf(codes.InvalidArgument, "size: %v", err) + } + if n < 0 { + return nil, status.Error(codes.InvalidArgument, "size must be non-negative") + } + walk = min(n, walk) + } + var sum uint32 + for i := int64(0); i < walk; i += pageSize { + sum ^= uint32(arr[i]) + } + s.ramReadBytes.Add(ctx, walk) + return &gluttonpb.ReadRAMResponse{Size: walk, Checksum: sum}, nil +} + +// Write to disk using the specified mode. Data written will be random bytes. +func (s *Service) WriteDisk(ctx context.Context, req *gluttonpb.WriteDiskRequest) (*gluttonpb.WriteDiskResponse, error) { + if !diskKeyRE.MatchString(req.GetKey()) { + return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) + } + if req.GetSize() < 0 { + return nil, status.Error(codes.InvalidArgument, "size must be non-negative") + } + + path := filepath.Join(s.dataDir, req.GetKey()) + + var flag int + switch req.GetWriteMode() { + case gluttonpb.WriteMode_WRITE_MODE_TRUNCATE: + flag = os.O_RDWR | os.O_CREATE | os.O_TRUNC + case gluttonpb.WriteMode_WRITE_MODE_OVERWRITE: + // No O_TRUNC: writes go from offset 0 but any bytes beyond size remain. + flag = os.O_RDWR | os.O_CREATE + default: + return nil, status.Errorf(codes.InvalidArgument, "unknown write_mode %v", req.GetWriteMode()) + } + + f, err := os.OpenFile(path, flag, 0o600) + if err != nil { + return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) + } + defer f.Close() + + h := sha256.New() + size := int64(req.GetSize()) + if err := streamRandomBytes(io.MultiWriter(f, h), size); err != nil { + return nil, status.Errorf(codes.Internal, "write %s: %v", path, err) + } + + // OVERWRITE has no O_TRUNC, bytes from a larger, earlier write will persist. + // The cursor is already at size, so folding the remainder into the + // same digest completes it without re-reading the prefix. + if req.GetWriteMode() == gluttonpb.WriteMode_WRITE_MODE_OVERWRITE { + tail, err := io.Copy(h, f) + if err != nil { + return nil, status.Errorf(codes.Internal, "hash tail %s: %v", path, err) + } + size += tail + } + + s.diskWriteBytes.Add(ctx, int64(req.GetSize())) + return &gluttonpb.WriteDiskResponse{Size: size, Sha256: h.Sum(nil)}, nil +} + +func (s *Service) ReadDisk(ctx context.Context, req *gluttonpb.ReadDiskRequest) (*gluttonpb.ReadDiskResponse, error) { + if !diskKeyRE.MatchString(req.GetKey()) { + return nil, status.Errorf(codes.InvalidArgument, "key %q must match %s", req.GetKey(), diskKeyRE) + } + + path := filepath.Join(s.dataDir, req.GetKey()) + + f, err := os.Open(path) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return nil, status.Errorf(codes.NotFound, "file %q not found", req.GetKey()) + } + return nil, status.Errorf(codes.Internal, "open %s: %v", path, err) + } + defer f.Close() + + h := sha256.New() + + if req.GetReadMode() == gluttonpb.ReadMode_READ_MODE_DIGEST_ONLY { + n, err := io.Copy(h, f) + if err != nil { + return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) + } + s.diskReadBytes.Add(ctx, n) + return &gluttonpb.ReadDiskResponse{ + Size: n, + Sha256: h.Sum(nil), + }, nil + } + + data, err := io.ReadAll(io.TeeReader(f, h)) + if err != nil { + return nil, status.Errorf(codes.Internal, "read %s: %v", path, err) + } + + s.diskReadBytes.Add(ctx, int64(len(data))) + return &gluttonpb.ReadDiskResponse{ + Size: int64(len(data)), + Sha256: h.Sum(nil), + Data: data, + }, nil +} + +// Make sure it has the specified number of file descriptors open. It will open or +// close file descriptors to hit the desired count (note this count is in addition to the other +// FDs needed to run the process). +func (s *Service) OpenFD(_ context.Context, req *gluttonpb.OpenFDRequest) (*gluttonpb.OpenFDResponse, error) { + if req.GetCount() < 0 { + return nil, status.Error(codes.InvalidArgument, "count must be non-negative") + } + target := int(req.GetCount()) + + s.mu.Lock() + defer s.mu.Unlock() + + for len(s.fds) > target { + last := len(s.fds) - 1 + if err := s.fds[last].Close(); err != nil { + slog.Warn("Failed to close glutton fd", slog.Any("err", err)) + } + s.fds[last] = nil + s.fds = s.fds[:last] + } + for len(s.fds) < target { + f, err := os.Open(os.DevNull) + if err != nil { + return nil, status.Errorf(codes.Internal, "open %s: %v", os.DevNull, err) + } + s.fds = append(s.fds, f) + } + return &gluttonpb.OpenFDResponse{}, nil +} + +// Receive a ping request, echoing the same response back. +func (s *Service) Ping(ctx context.Context, req *gluttonpb.PingRequest) (*gluttonpb.PingResponse, error) { + s.pingsReceived.Add(ctx, 1) + return &gluttonpb.PingResponse{Message: req.GetMessage()}, nil +} + +func randomBytes(n int) ([]byte, error) { + buf := make([]byte, n) + if _, err := rand.Read(buf); err != nil { + return nil, err + } + return buf, nil +} + +// streamRandomBytesChunk caps per-syscall random fill and write size so a +// multi-gigabyte WriteDisk doesn't have to materialize in RAM. +const streamRandomBytesChunk = 1 << 20 // 1 MiB + +// streamRandomBytes writes total random bytes to w sequentially, in +// streamRandomBytesChunk-sized chunks. The caller is responsible for the +// file's open mode and starting offset; this writes from the current +// position forward. +func streamRandomBytes(w io.Writer, total int64) error { + if total <= 0 { + return nil + } + buf := make([]byte, streamRandomBytesChunk) + var written int64 + for written < total { + chunk := buf + if remaining := total - written; remaining < int64(len(chunk)) { + chunk = buf[:remaining] + } + if _, err := rand.Read(chunk); err != nil { + return fmt.Errorf("generate random bytes: %w", err) + } + n, err := w.Write(chunk) + if err != nil { + return err + } + written += int64(n) + } + return nil +} diff --git a/internal/benchmarking/glutton/service_test.go b/internal/benchmarking/glutton/service_test.go new file mode 100644 index 0000000000..f2e1fbc6de --- /dev/null +++ b/internal/benchmarking/glutton/service_test.go @@ -0,0 +1,226 @@ +// Copyright 2026 Google LLC +// +// 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 glutton + +import ( + "bytes" + "context" + "crypto/sha256" + "os" + "path/filepath" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + gluttonpb "github.com/agent-substrate/substrate/internal/proto/glutton" +) + +func TestWriteDiskReadDiskRoundTrip(t *testing.T) { + tempDir := t.TempDir() + svc, err := New(tempDir) + if err != nil { + t.Fatalf("failed to create glutton service: %v", err) + } + defer svc.Close() + + ctx := context.Background() + tests := []struct { + name string + key string + size int32 + }{ + {name: "zero size", key: "zero", size: 0}, + {name: "small size", key: "small", size: 1024}, + {name: "chunk unaligned size", key: "unaligned", size: (1 << 20) + 1}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + writeResp, err := svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ + Key: tt.key, + Size: tt.size, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, + }) + if err != nil { + t.Fatalf("WriteDisk failed: %v", err) + } + if writeResp.GetSize() != int64(tt.size) { + t.Errorf("WriteDisk size mismatch: got %d, want %d", writeResp.GetSize(), tt.size) + } + + // 1. Full data read + readResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ + Key: tt.key, + ReadMode: gluttonpb.ReadMode_READ_MODE_DATA, + }) + if err != nil { + t.Fatalf("ReadDisk (DATA) failed: %v", err) + } + + if readResp.GetSize() != int64(tt.size) { + t.Errorf("ReadDisk size mismatch: got %d, want %d", readResp.GetSize(), tt.size) + } + if !bytes.Equal(readResp.GetSha256(), writeResp.GetSha256()) { + t.Errorf("sha256 mismatch between WriteDisk and ReadDisk") + } + if len(readResp.GetData()) != int(tt.size) { + t.Errorf("ReadDisk data length mismatch: got %d, want %d", len(readResp.GetData()), tt.size) + } + + computedDigest := sha256.Sum256(readResp.GetData()) + if !bytes.Equal(readResp.GetSha256(), computedDigest[:]) { + t.Errorf("ReadDisk returned sha256 does not match computed digest of returned data") + } + + // 2. Digest-only read + digestResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ + Key: tt.key, + ReadMode: gluttonpb.ReadMode_READ_MODE_DIGEST_ONLY, + }) + if err != nil { + t.Fatalf("ReadDisk (DIGEST_ONLY) failed: %v", err) + } + if digestResp.GetSize() != int64(tt.size) { + t.Errorf("ReadDisk (DIGEST_ONLY) size mismatch: got %d, want %d", digestResp.GetSize(), tt.size) + } + if !bytes.Equal(digestResp.GetSha256(), writeResp.GetSha256()) { + t.Errorf("sha256 mismatch between WriteDisk and ReadDisk (DIGEST_ONLY)") + } + if len(digestResp.GetData()) != 0 { + t.Errorf("ReadDisk (DIGEST_ONLY) should not return data payload, got %d bytes", len(digestResp.GetData())) + } + }) + } +} + +func TestWriteDiskTruncateProducesExactSize(t *testing.T) { + tempDir := t.TempDir() + svc, err := New(tempDir) + if err != nil { + t.Fatalf("failed to create glutton service: %v", err) + } + defer svc.Close() + + ctx := context.Background() + key := "testfile" + size := int32(2048) + + _, err = svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ + Key: key, + Size: size, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, + }) + if err != nil { + t.Fatalf("WriteDisk failed: %v", err) + } + + filePath := filepath.Join(tempDir, key) + fi, err := os.Stat(filePath) + if err != nil { + t.Fatalf("os.Stat failed: %v", err) + } + if fi.Size() != int64(size) { + t.Errorf("file size on disk mismatch: got %d, want %d", fi.Size(), size) + } +} + +func TestWriteDiskOverwriteDigestMatchesReadDisk(t *testing.T) { + tempDir := t.TempDir() + svc, err := New(tempDir) + if err != nil { + t.Fatalf("failed to create glutton service: %v", err) + } + defer svc.Close() + + ctx := context.Background() + key := "overwrittenfile" + + // 1. Initial write of large file (4096 bytes) + _, err = svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ + Key: key, + Size: 4096, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_TRUNCATE, + }) + if err != nil { + t.Fatalf("WriteDisk (large) failed: %v", err) + } + + // 2. Overwrite prefix with smaller size (1024 bytes) without truncation + overwriteResp, err := svc.WriteDisk(ctx, &gluttonpb.WriteDiskRequest{ + Key: key, + Size: 1024, + WriteMode: gluttonpb.WriteMode_WRITE_MODE_OVERWRITE, + }) + if err != nil { + t.Fatalf("WriteDisk (overwrite) failed: %v", err) + } + + if overwriteResp.GetSize() != 4096 { + t.Errorf("expected WriteDisk under OVERWRITE to report total file size 4096, got %d", overwriteResp.GetSize()) + } + + // 3. ReadDisk reads the entire file (4096 bytes) + readResp, err := svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{ + Key: key, + ReadMode: gluttonpb.ReadMode_READ_MODE_DATA, + }) + if err != nil { + t.Fatalf("ReadDisk failed: %v", err) + } + + if readResp.GetSize() != 4096 { + t.Errorf("expected ReadDisk size 4096, got %d", readResp.GetSize()) + } + if !bytes.Equal(readResp.GetSha256(), overwriteResp.GetSha256()) { + t.Errorf("expected WriteDisk(OVERWRITE) whole-file digest to match ReadDisk digest") + } +} + +func TestReadDiskRejectsInvalidKey(t *testing.T) { + tempDir := t.TempDir() + svc, err := New(tempDir) + if err != nil { + t.Fatalf("failed to create glutton service: %v", err) + } + defer svc.Close() + + ctx := context.Background() + _, err = svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{Key: "../escape"}) + if err == nil { + t.Error("expected error for invalid key with path traversal, got nil") + } + if s, ok := status.FromError(err); !ok || s.Code() != codes.InvalidArgument { + t.Errorf("expected InvalidArgument code, got %v", err) + } +} + +func TestReadDiskNotFound(t *testing.T) { + tempDir := t.TempDir() + svc, err := New(tempDir) + if err != nil { + t.Fatalf("failed to create glutton service: %v", err) + } + defer svc.Close() + + ctx := context.Background() + _, err = svc.ReadDisk(ctx, &gluttonpb.ReadDiskRequest{Key: "nonexistent"}) + if err == nil { + t.Error("expected error for nonexistent file, got nil") + } + if s, ok := status.FromError(err); !ok || s.Code() != codes.NotFound { + t.Errorf("expected NotFound code, got %v", err) + } +}