feat: add explicit SUBSCRIBE message to initiate subscriptions

Author: jbedard

runner/pkg/ibp/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/socket",
+        "@aspect_gazelle//common/logger",
         "@com_github_fatih_color//:color",
     ],
 )

runner/pkg/ibp/client.go

@@ -5,23 +5,27 @@ import (
 	"iter"
 	"slices"
 
+	BazelLog "github.com/aspect-build/aspect-gazelle/common/logger"
 	"github.com/aspect-build/aspect-gazelle/runner/pkg/socket"
 )
 
-type IncrementalClient interface {
+type WatchClient interface {
 	Connect() error
 	Disconnect() error
-	AwaitCycle() iter.Seq[CycleSourcesMessage]
+	Subscribe(options WatchOptions) iter.Seq2[*CycleSourcesMessage, error]
 }
 
 type incClient struct {
 	socketPath string
 	socket     socket.Socket[interface{}, map[string]interface{}]
+
+	// The negotiated protocol version
+	version ProtocolVersion
 }
 
-var _ IncrementalClient = (*incClient)(nil)
+var _ WatchClient = (*incClient)(nil)
 
-func NewClient(host string) IncrementalClient {
+func NewClient(host string) WatchClient {
 	return &incClient{
 		socketPath: host,
 	}
@@ -55,22 +59,39 @@ func (c *incClient) negotiate() error {
 	if negReq["versions"] == nil {
 		return fmt.Errorf("Received NEGOTIATE without versions: %v", negReq)
 	}
-	if !slices.Contains(negReq["versions"].([]interface{}), (interface{})(float64(PROTOCOL_VERSION))) {
-		return fmt.Errorf("Received NEGOTIATE with unsupported versions %v, expected %d", negReq["versions"], PROTOCOL_VERSION)
+	rawVersions, isArray := negReq["versions"].([]interface{})
+	if !isArray {
+		return fmt.Errorf("Invalid versions, expected []int, received type: %T", negReq["versions"])
+	}
+
+	negotiatedVersion, err := negoiateVersion(rawVersions)
+	if err != nil {
+		return err
 	}
 
 	err = c.socket.Send(negotiateResponseMessage{
 		Message: Message{
 			Kind: "NEGOTIATE_RESPONSE",
 		},
-		Version: PROTOCOL_VERSION,
+		Version: negotiatedVersion,
 	})
 	if err != nil {
 		return fmt.Errorf("failed to negotiate protocol version: %w", err)
 	}
+
+	c.version = negotiatedVersion
 	return nil
 }
 
+func negoiateVersion(acceptedVersions []interface{}) (ProtocolVersion, error) {
+	for _, v := range slices.Backward(acceptedVersions) {
+		if slices.Contains(abazelSupportedProtocolVersions, ProtocolVersion(v.(float64))) {
+			return ProtocolVersion(v.(float64)), nil
+		}
+	}
+	return -1, fmt.Errorf("unsupported versions %v, expected one of %v", acceptedVersions, abazelSupportedProtocolVersions)
+}
+
 func (c *incClient) Disconnect() error {
 	if c.socket == nil {
 		return fmt.Errorf("client not connected")
@@ -84,12 +105,37 @@ func (c *incClient) Disconnect() error {
 	return err
 }
 
-func (c *incClient) AwaitCycle() iter.Seq[CycleSourcesMessage] {
-	return func(yield func(CycleSourcesMessage) bool) {
+func (c *incClient) Subscribe(options WatchOptions) iter.Seq2[*CycleSourcesMessage, error] {
+	return func(yield func(*CycleSourcesMessage, error) bool) {
+		if c.version >= VERSION_1 {
+			err := c.socket.Send(SubscribeMessage{
+				Message: Message{
+					Kind: "SUBSCRIBE",
+				},
+				WatchType: options.Type,
+			})
+			if err != nil {
+				yield(nil, err)
+				return
+			}
+
+			msg, err := c.socket.Recv()
+			if err != nil {
+				yield(nil, err)
+				return
+			}
+
+			if msg["kind"] != "SUBSCRIBE_RESPONSE" {
+				yield(nil, fmt.Errorf("expected SUBSCRIBE_RESPONSE, got %v", msg))
+				return
+			}
+		}
+
 		for {
 			msg, err := c.socket.Recv()
 			if err != nil {
 				fmt.Printf("Error receiving message: %v\n", err)
+				yield(nil, err)
 				return
 			}
 
@@ -100,21 +146,28 @@ func (c *incClient) AwaitCycle() iter.Seq[CycleSourcesMessage] {
 					continue
 				}
 
-				c.socket.Send(CycleMessage{
+				err = c.socket.Send(CycleMessage{
 					Message: Message{
 						Kind: "CYCLE_STARTED",
 					},
 					CycleId: cycleEvent.CycleId,
 				})
+				if err != nil {
+					yield(nil, err)
+					return
+				}
 
-				r := yield(cycleEvent)
+				r := yield(&cycleEvent, nil)
 
-				c.socket.Send(CycleMessage{
+				err = c.socket.Send(CycleMessage{
 					Message: Message{
 						Kind: "CYCLE_COMPLETED",
 					},
 					CycleId: cycleEvent.CycleId,
 				})
+				if err != nil {
+					BazelLog.Warnf("Failed to send CYCLE_COMPLETED for cycle_id=%d: %v\n", cycleEvent.CycleId, err)
+				}
 
 				if !r {
 					return

runner/pkg/ibp/protocol.go

@@ -5,13 +5,25 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"slices"
 	"sync/atomic"
 
 	"github.com/aspect-build/aspect-gazelle/runner/pkg/socket"
 	"github.com/fatih/color"
 )
 
-const PROTOCOL_VERSION = 0
+type ProtocolVersion int
+
+const (
+	LEGACY_VERSION_0 ProtocolVersion = 0
+	VERSION_1                        = 1
+	LATEST_VERSION                   = VERSION_1
+)
+
+func (v ProtocolVersion) HasExplicitSubscribe() bool {
+	return v >= 1
+}
+
 const PROTOCOL_SOCKET_ENV = "ABAZEL_WATCH_SOCKET_FILE"
 
 type IncrementalBazel interface {
@@ -40,18 +52,34 @@ type Message struct {
 
 type negotiateMessage struct {
 	Message
-	Versions []int `json:"versions"`
+	Versions []ProtocolVersion `json:"versions"`
 }
 type negotiateResponseMessage struct {
 	Message
-	Version int `json:"version"`
+	Version ProtocolVersion `json:"version"`
 }
 
 type capMessage struct {
 	Message
 	Caps map[string]bool `json:"caps"`
 }
 
+type WatchType string
+
+const (
+	WatchTypeRunfiles WatchType = "runfiles"
+	WatchTypeSources  WatchType = "sources"
+)
+
+type WatchOptions struct {
+	Type WatchType
+}
+
+type SubscribeMessage struct {
+	Message
+	WatchType WatchType `json:"watch_type"`
+}
+
 type exitMessage struct {
 	Message
 	Description string `json:"description"`
@@ -77,7 +105,10 @@ type CycleSourcesMessage struct {
 }
 
 // The versions supported by this host implementation of the protocol.
-var abazelSupportedProtocolVersions = []int{PROTOCOL_VERSION}
+var abazelSupportedProtocolVersions = []ProtocolVersion{
+	LEGACY_VERSION_0,
+	VERSION_1,
+}
 
 type aspectBazelSocket = socket.Server[interface{}, map[string]any]
 
@@ -172,8 +203,8 @@ func (p *aspectBazelProtocol) acceptNegotiation() error {
 	if negResp["version"] == nil {
 		return fmt.Errorf("Received NEGOTIATE_RESPONSE without version: %v", negResp)
 	}
-	if negResp["version"].(float64) != PROTOCOL_VERSION {
-		return fmt.Errorf("Received NEGOTIATE_RESPONSE with unsupported version %v, expected %d", negResp["version"], PROTOCOL_VERSION)
+	if !slices.Contains(abazelSupportedProtocolVersions, ProtocolVersion(negResp["version"].(float64))) {
+		return fmt.Errorf("Received NEGOTIATE_RESPONSE with unsupported version %v, expected one of %v", negResp["version"], abazelSupportedProtocolVersions)
 	}
 
 	p.connectedCh <- int(negResp["version"].(float64))

runner/runner.go

@@ -238,8 +238,8 @@ func (p *GazelleRunner) Watch(watchAddress string, cmd GazelleCommand, mode Gaze
 	))
 	defer t.End()
 
-	// Subscribe to further changes
-	for cs := range watch.AwaitCycle() {
+	// Subscribe to further changes to all sources
+	for cs := range watch.Subscribe(ibp.WatchOptions{Type: ibp.WatchTypeSources}) {
 		_, t := p.tracer.Start(ctx, "GazelleRunner.Watch.Trigger")
 
 		// The directories that have changed which gazelle should update.