internal/coordinator/remote, internal/gomote: add instance alive

This change adds the implementation for gomote instance alive. This
endpoint enables the caller to check if an instance is alive. If the
instance is alive, it will extend the gomote timeout time. A renew
timeout method has been added to the session pool. This differs from
the existing keep alive method in that it is an single call to renew
the timeout for an instance instead of a continuous renewal of the
timeout tied to the lifetime of a context.

Updated golang/go#48742

Change-Id: I3b3462407d9f4a02c4e2cea0f14950c8c9f21060
Reviewed-on: https://go-review.googlesource.com/c/build/+/374115
Trust: Carlos Amedee <[email protected]>
Run-TryBot: Carlos Amedee <[email protected]>
Reviewed-by: Dmitri Shuralyov <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>

Author: cagedmantis

internal/coordinator/remote/remote.go

@@ -228,3 +228,16 @@ func (sp *SessionPool) KeepAlive(ctx context.Context, buildletName string) error
 	})
 	return nil
 }
+
+// RenewTimeout will renew the remote buildlet session by extending the expiration value.
+func (sp *SessionPool) RenewTimeout(buildletName string) error {
+	sp.mu.Lock()
+	defer sp.mu.Unlock()
+
+	s, ok := sp.m[buildletName]
+	if !ok {
+		return fmt.Errorf("remote buildlet does not exist=%s", buildletName)
+	}
+	s.renew()
+	return nil
+}

internal/coordinator/remote/remote_test.go

@@ -95,3 +95,23 @@ func TestSessionPoolDestroySession(t *testing.T) {
 		}
 	}
 }
+
+func TestRenewTimeout(t *testing.T) {
+	sp := NewSessionPool(context.Background())
+	defer sp.Close()
+
+	name := sp.AddSession("accounts.google.com:user-xyz-124", "user-x", "builder", "host", &buildlet.FakeClient{})
+	if err := sp.RenewTimeout(name); err != nil {
+		t.Errorf("SessionPool.RenewTimeout(%q) = %s; want no error", name, err)
+	}
+}
+
+func TestRenewTimeoutError(t *testing.T) {
+	sp := NewSessionPool(context.Background())
+	defer sp.Close()
+
+	name := sp.AddSession("accounts.google.com:user-xyz-124", "user-x", "builder", "host", &buildlet.FakeClient{})
+	if err := sp.RenewTimeout(name + "-wrong"); err == nil {
+		t.Errorf("SessionPool.RenewTimeout(%q) = %s; want error", name, err)
+	}
+}

internal/gomote/gomote.go

@@ -139,6 +139,29 @@ func (s *Server) CreateInstance(req *protos.CreateInstanceRequest, stream protos
 	}
 }
 
+// InstanceAlive will ensure that the gomote instance is still alive and will extend the timeout. The requester must be authenticated.
+func (s *Server) InstanceAlive(ctx context.Context, req *protos.InstanceAliveRequest) (*protos.InstanceAliveResponse, error) {
+	creds, err := access.IAPFromContext(ctx)
+	if err != nil {
+		log.Printf("InstanceAlive access.IAPFromContext(ctx) = nil, %s", err)
+		return nil, status.Errorf(codes.Unauthenticated, "request does not contain the required authentication")
+	}
+	if req.GetGomoteId() == "" {
+		return nil, status.Errorf(codes.InvalidArgument, "invalid gomote ID")
+	}
+	session, err := s.buildlets.Session(req.GetGomoteId())
+	if err != nil {
+		return nil, status.Errorf(codes.NotFound, "specified gomote instance does not exist")
+	}
+	if session.OwnerID != creds.ID {
+		return nil, status.Errorf(codes.PermissionDenied, "not allowed to modify this gomote session")
+	}
+	if err := s.buildlets.RenewTimeout(req.GetGomoteId()); err != nil {
+		return nil, status.Errorf(codes.Internal, "unable to renew timeout")
+	}
+	return &protos.InstanceAliveResponse{}, nil
+}
+
 // ListInstances will list the gomote instances owned by the requester. The requester must be authenticated.
 func (s *Server) ListInstances(ctx context.Context, req *protos.ListInstancesRequest) (*protos.ListInstancesResponse, error) {
 	creds, err := access.IAPFromContext(ctx)

internal/gomote/gomote_test.go

@@ -157,6 +157,75 @@ func TestCreateInstanceError(t *testing.T) {
 	}
 }
 
+func TestInstanceAlive(t *testing.T) {
+	client := setupGomoteTest(t, context.Background())
+	gomoteID := mustCreateInstance(t, client, fakeIAP())
+	req := &protos.InstanceAliveRequest{
+		GomoteId: gomoteID,
+	}
+	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())
+	got, err := client.InstanceAlive(ctx, req)
+	if err != nil {
+		t.Fatalf("client.InstanceAlive(ctx, %v) = %v, %s; want no error", req, got, err)
+	}
+}
+
+func TestInstanceAliveError(t *testing.T) {
+	// This test will create a gomote instance and attempt to call InstanceAlive.
+	// If overrideID is set to true, the test will use a diffrent gomoteID than the
+	// the one created for the test.
+	testCases := []struct {
+		desc       string
+		ctx        context.Context
+		overrideID bool
+		gomoteID   string // Used iff overrideID is true.
+		wantCode   codes.Code
+	}{
+		{
+			desc:     "unauthenticated request",
+			ctx:      context.Background(),
+			wantCode: codes.Unauthenticated,
+		},
+		{
+			desc:       "missing gomote id",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
+			overrideID: true,
+			wantCode:   codes.InvalidArgument,
+		},
+		{
+			desc:       "gomote does not exist",
+			ctx:        access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP()),
+			overrideID: true,
+			gomoteID:   "xyz",
+			wantCode:   codes.NotFound,
+		},
+		{
+			desc:     "gomote is not owned by caller",
+			ctx:      access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAPWithUser("user-x", "email-y")),
+			wantCode: codes.PermissionDenied,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			client := setupGomoteTest(t, context.Background())
+			gomoteID := mustCreateInstance(t, client, fakeIAP())
+			if tc.overrideID {
+				gomoteID = tc.gomoteID
+			}
+			req := &protos.InstanceAliveRequest{
+				GomoteId: gomoteID,
+			}
+			got, err := client.InstanceAlive(tc.ctx, req)
+			if err != nil && status.Code(err) != tc.wantCode {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if err == nil {
+				t.Fatalf("client.InstanceAlive(ctx, %v) = %v, nil; want error", req, got)
+			}
+		})
+	}
+}
+
 func TestListInstance(t *testing.T) {
 	client := setupGomoteTest(t, context.Background())
 	ctx := access.FakeContextWithOutgoingIAPAuth(context.Background(), fakeIAP())