cmd/coordinator: check the session pool for buildlets

Remote buildlets should not be destroyed by the pool implementations for
instances which have been created on clouds. This change adds a check
to ensure that the pool implementations can check for remote buildlets
for instances which are managed by the session pools. Those sessions will
be destroyed by the session pool.

The buildlet client accessors for instance names have been changed so
that they are no longer GCE specific.

The EC2 session pool implementation has been changed so that the
instance name is set in the buildlet client when an EC2 instance is
created.

Updates golang/go#54735
Updates golang/go#54696

Change-Id: Iae3c86cc47b4b79e3287d581dc042dcf2714f7c9
Reviewed-on: https://go-review.googlesource.com/c/build/+/426015
Reviewed-by: Carlos Amedee <carlos@golang.org>
Auto-Submit: Carlos Amedee <carlos@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Run-TryBot: Carlos Amedee <carlos@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>

buildlet/buildletclient.go

@@ -108,17 +108,17 @@ func (c *client) SetDescription(v string) {
 	c.desc = v
 }
 
-// SetGCEInstanceName sets an instance name for GCE buildlets.
+// SetInstanceName sets an instance name for GCE and EC2 buildlets.
 // This value differs from the buildlet name used in the CLI and web interface.
-func (c *client) SetGCEInstanceName(v string) {
-	c.gceInstanceName = v
+func (c *client) SetInstanceName(v string) {
+	c.instanceName = v
 }
 
-// GCEInstanceName gets an instance name for GCE buildlets.
+// InstanceName gets an instance name for GCE and EC2 buildlets.
 // This value differs from the buildlet name used in the CLI and web interface.
-// For non-GCE buildlets, this will return an empty string.
-func (c *client) GCEInstanceName() string {
-	return c.gceInstanceName
+// For non-GCE or EC2 buildlets, this will return an empty string.
+func (c *client) InstanceName() string {
+	return c.instanceName
 }
 
 // SetHTTPClient replaces the underlying HTTP client.
@@ -153,16 +153,16 @@ func defaultDialer() func(network, addr string) (net.Conn, error) {
 
 // A client interacts with a single buildlet.
 type client struct {
-	ipPort          string // required, unless remoteBuildlet+baseURL is set
-	tls             KeyPair
-	httpClient      *http.Client
-	dialer          func(context.Context) (net.Conn, error) // nil means to use net.Dialer.DialContext
-	baseURL         string                                  // optional baseURL (used by remote buildlets)
-	authUser        string                                  // defaults to "gomote", if password is non-empty
-	password        string                                  // basic auth password or empty for none
-	remoteBuildlet  string                                  // non-empty if for remote buildlets (used by client)
-	name            string                                  // optional name for debugging, returned by Name
-	gceInstanceName string                                  // instance name for GCE VMs
+	ipPort         string // required, unless remoteBuildlet+baseURL is set
+	tls            KeyPair
+	httpClient     *http.Client
+	dialer         func(context.Context) (net.Conn, error) // nil means to use net.Dialer.DialContext
+	baseURL        string                                  // optional baseURL (used by remote buildlets)
+	authUser       string                                  // defaults to "gomote", if password is non-empty
+	password       string                                  // basic auth password or empty for none
+	remoteBuildlet string                                  // non-empty if for remote buildlets (used by client)
+	name           string                                  // optional name for debugging, returned by Name
+	instanceName   string                                  // instance name for GCE and EC2 VMs
 
 	closeFuncs []func() // optional extra code to run on close
 

buildlet/fakebuildletclient.go

@@ -38,16 +38,16 @@ type RemoteClient interface {
 type Client interface {
 	RemoteClient
 	ConnectSSH(user, authorizedPubKey string) (net.Conn, error)
-	GCEInstanceName() string
 	IPPort() string
+	InstanceName() string
 	IsBroken() bool
 	MarkBroken()
 	Name() string
 	ProxyRoundTripper() http.RoundTripper
 	SetDescription(v string)
 	SetDialer(dialer func(context.Context) (net.Conn, error))
-	SetGCEInstanceName(v string)
 	SetHTTPClient(httpClient *http.Client)
+	SetInstanceName(v string)
 	SetName(name string)
 	SetOnHeartbeatFailure(fn func())
 	Status(ctx context.Context) (Status, error)
@@ -96,8 +96,8 @@ func (fc *FakeClient) Exec(ctx context.Context, cmd string, opts ExecOpts) (remo
 	return nil, nil
 }
 
-// GCEInstanceName gives the fake instance name.
-func (fc *FakeClient) GCEInstanceName() string { return fc.instanceName }
+// InstanceName gives the fake instance name.
+func (fc *FakeClient) InstanceName() string { return fc.instanceName }
 
 // GetTar gives a vake tar zipped directory.
 func (fc *FakeClient) GetTar(ctx context.Context, dir string) (io.ReadCloser, error) {
@@ -166,8 +166,8 @@ func (fc *FakeClient) SetDescription(v string) {}
 // SetDialer sets the function that creates a new connection to the fake buildlet.
 func (fc *FakeClient) SetDialer(dialer func(context.Context) (net.Conn, error)) {}
 
-// SetGCEInstanceName sets the GCE instance name on a fake client.
-func (fc *FakeClient) SetGCEInstanceName(v string) {
+// SetInstanceName sets the GCE or EC2 instance name on a fake client.
+func (fc *FakeClient) SetInstanceName(v string) {
 	fc.instanceName = v
 }
 

cmd/coordinator/coordinator.go

@@ -249,7 +249,13 @@ func main() {
 	// a shared package.
 	pool.SetBuilderMasterKey(masterKey())
 
-	err := pool.InitGCE(sc, &basePinErr, isGCERemoteBuildlet, *buildEnvName, *mode)
+	sp := remote.NewSessionPool(context.Background())
+	// TODO(go.dev/issue/54735): remove while legacy gomote implementation is being removed.
+	isRemoteBuildlet := func(instName string) bool {
+		return isGCERemoteBuildlet(instName) || sp.IsSession(instName)
+	}
+
+	err := pool.InitGCE(sc, &basePinErr, isRemoteBuildlet, *buildEnvName, *mode)
 	if err != nil {
 		if *mode == "" {
 			*mode = "dev"
@@ -277,7 +283,7 @@ func main() {
 	if *mode == "prod" || (*mode == "dev" && *devEnableEC2) {
 		// TODO(golang.org/issues/38337) the coordinator will use a package scoped pool
 		// until the coordinator is refactored to not require them.
-		ec2Pool := mustCreateEC2BuildletPool(sc)
+		ec2Pool := mustCreateEC2BuildletPool(sc, isRemoteBuildlet)
 		defer ec2Pool.Close()
 	}
 
@@ -360,7 +366,6 @@ func main() {
 	dashV1 := legacydash.Handler(gce.GoDSClient(), maintnerClient, string(masterKey()), grpcServer)
 	dashV2 := &builddash.Handler{Datastore: gce.GoDSClient(), Maintner: maintnerClient}
 	gs := &gRPCServer{dashboardURL: "https://build.golang.org"}
-	sp := remote.NewSessionPool(context.Background())
 	setSessionPool(sp)
 	gomoteServer := gomote.New(sp, sched, sshCA, gomoteBucket, mustStorageClient())
 	protos.RegisterCoordinatorServer(grpcServer, gs)
@@ -2204,7 +2209,7 @@ func mustCreateSecretClientOnGCE() *secret.Client {
 	return secret.MustNewClient()
 }
 
-func mustCreateEC2BuildletPool(sc *secret.Client) *pool.EC2Buildlet {
+func mustCreateEC2BuildletPool(sc *secret.Client, isRemoteBuildlet func(instName string) bool) *pool.EC2Buildlet {
 	awsKeyID, err := sc.Retrieve(context.Background(), secret.NameAWSKeyID)
 	if err != nil {
 		log.Fatalf("unable to retrieve secret %q: %s", secret.NameAWSKeyID, err)
@@ -2220,7 +2225,7 @@ func mustCreateEC2BuildletPool(sc *secret.Client) *pool.EC2Buildlet {
 		log.Fatalf("unable to create AWS client: %s", err)
 	}
 
-	ec2Pool, err := pool.NewEC2Buildlet(awsClient, buildenv.Production, dashboard.Hosts, isGCERemoteBuildlet)
+	ec2Pool, err := pool.NewEC2Buildlet(awsClient, buildenv.Production, dashboard.Hosts, isRemoteBuildlet)
 	if err != nil {
 		log.Fatalf("unable to create EC2 buildlet pool: %s", err)
 	}

cmd/coordinator/remote.go

@@ -69,7 +69,7 @@ func isGCERemoteBuildlet(instName string) bool {
 	remoteBuildlets.Lock()
 	defer remoteBuildlets.Unlock()
 	for _, rb := range remoteBuildlets.M {
-		if rb.Buildlet().GCEInstanceName() == instName {
+		if rb.Buildlet().InstanceName() == instName {
 			return true
 		}
 	}

internal/coordinator/pool/ec2.go

@@ -219,6 +219,7 @@ func (eb *EC2Buildlet) GetBuildlet(ctx context.Context, hostType string, lg Logg
 		log.Printf("EC2 VM %q failed heartbeat", instName)
 		eb.buildletDone(instName)
 	})
+	bc.SetInstanceName(instName)
 	return bc, nil
 }
 

internal/coordinator/pool/gce.go

@@ -82,16 +82,16 @@ var (
 	gkeNodeHostname string
 
 	// values created due to separating the buildlet pools into a separate package
-	gceMode             string
-	basePinErr          *atomic.Value
-	isGCERemoteBuildlet IsRemoteBuildletFunc
+	gceMode          string
+	basePinErr       *atomic.Value
+	isRemoteBuildlet IsRemoteBuildletFunc
 )
 
 // InitGCE initializes the GCE buildlet pool.
 func InitGCE(sc *secret.Client, basePin *atomic.Value, fn IsRemoteBuildletFunc, buildEnvName, mode string) error {
 	gceMode = mode
 	basePinErr = basePin
-	isGCERemoteBuildlet = fn
+	isRemoteBuildlet = fn
 
 	ctx := context.Background()
 	var err error
@@ -495,7 +495,7 @@ func (p *GCEBuildlet) GetBuildlet(ctx context.Context, hostType string, lg Logge
 	}
 	waitBuildlet.Done(nil)
 	bc.SetDescription("GCE VM: " + instName)
-	bc.SetGCEInstanceName(instName)
+	bc.SetInstanceName(instName)
 	bc.SetOnHeartbeatFailure(cleanup)
 	return bc, nil
 }
@@ -650,7 +650,7 @@ func (p *GCEBuildlet) cleanZoneVMs(zone string) error {
 				// Defensive. Not seen in practice.
 				continue
 			}
-			if isGCERemoteBuildlet(inst.Name) {
+			if isRemoteBuildlet(inst.Name) {
 				// Remote buildlets have their own expiration mechanism that respects active SSH sessions.
 				log.Printf("cleanZoneVMs: skipping remote buildlet %q", inst.Name)
 				continue

internal/coordinator/remote/remote.go

@@ -96,13 +96,14 @@ func (sp *SessionPool) AddSession(ownerID, username, builderType, hostType strin
 	}
 }
 
-// IsGCESession checks if the session is a GCE instance.
-func (sp *SessionPool) IsGCESession(instName string) bool {
+// IsSession is true if the instance is found in the session pool. The instance name is the not the public
+// name of the instance. It is the name of the instance as it is tracked in the cloud service.
+func (sp *SessionPool) IsSession(instName string) bool {
 	sp.mu.RLock()
 	defer sp.mu.RUnlock()
 
 	for _, s := range sp.m {
-		if s.buildlet.GCEInstanceName() == instName {
+		if s.buildlet.InstanceName() == instName {
 			return true
 		}
 	}