diff --git a/cmd/benchmarking/glutton/main.go b/cmd/benchmarking/glutton/main.go index 5488402745..9c01fb99b6 100644 --- a/cmd/benchmarking/glutton/main.go +++ b/cmd/benchmarking/glutton/main.go @@ -19,49 +19,23 @@ 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.") 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.") ) @@ -81,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 { @@ -89,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) } @@ -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/cmd/benchmarking/glutton/main_test.go b/cmd/benchmarking/glutton/main_test.go deleted file mode 100644 index 34d8db2710..0000000000 --- a/cmd/benchmarking/glutton/main_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 main - -import ( - "bytes" - "context" - "crypto/sha256" - "io" - "net" - "net/http" - "net/http/httptest" - "os" - "path/filepath" - "testing" - "time" - - "github.com/agent-substrate/substrate/internal/ateinterceptors" - "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 := newGluttonService(t.TempDir()) - if err != nil { - t.Fatalf("newGluttonService: %v", err) - } - defer svc.Close() - - grpcSrv := grpc.NewServer() - glutton.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 := glutton.NewGluttonClient(conn).Ping(ctx, &glutton.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 := newGluttonService(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, &glutton.WriteDiskRequest{ - Key: tt.key, - Size: tt.size, - WriteMode: glutton.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, &glutton.ReadDiskRequest{ - Key: tt.key, - ReadMode: glutton.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, &glutton.ReadDiskRequest{ - Key: tt.key, - ReadMode: glutton.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 := newGluttonService(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, &glutton.WriteDiskRequest{ - Key: key, - Size: size, - WriteMode: glutton.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 := newGluttonService(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, &glutton.WriteDiskRequest{ - Key: key, - Size: 4096, - WriteMode: glutton.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{ - Key: key, - Size: 1024, - WriteMode: glutton.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, &glutton.ReadDiskRequest{ - Key: key, - ReadMode: glutton.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 := newGluttonService(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"}) - 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 := newGluttonService(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"}) - 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 := newGluttonService(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(&glutton.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 glutton.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(&glutton.WriteDiskRequest{ - Key: "httpfile", - Size: 512, - WriteMode: glutton.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 glutton.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(&glutton.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 glutton.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(&glutton.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(&glutton.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/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 new file mode 100644 index 0000000000..aaa615f3d5 --- /dev/null +++ b/internal/benchmarking/glutton/glutton.go @@ -0,0 +1,44 @@ +// 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 + +// 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" +) + +// 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/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/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/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) + } +} 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") } }