| // Copyright 2020 The Go Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style |
| // license that can be found in the LICENSE file. |
| |
| //go:build linux || darwin |
| // +build linux darwin |
| |
| package pool |
| |
| import ( |
| "context" |
| "errors" |
| "fmt" |
| "sort" |
| "testing" |
| "time" |
| |
| "github.com/google/go-cmp/cmp" |
| "golang.org/x/build/buildenv" |
| "golang.org/x/build/buildlet" |
| "golang.org/x/build/dashboard" |
| "golang.org/x/build/internal/cloud" |
| "golang.org/x/build/internal/spanlog" |
| ) |
| |
| func TestEC2BuildletGetBuildlet(t *testing.T) { |
| host := "host-type-x" |
| |
| l := newLedger() |
| l.UpdateInstanceTypes([]*cloud.InstanceType{ |
| // set to default gce type because there is no way to set the machine |
| // type from outside of the buildenv package. |
| { |
| Type: "e2-standard-4", |
| CPU: 4, |
| }, |
| }) |
| l.SetCPULimit(20) |
| |
| bp := &EC2Buildlet{ |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: true, |
| VMCreated: true, |
| buildletCreated: true, |
| }, |
| buildEnv: &buildenv.Environment{}, |
| ledger: l, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: { |
| VMImage: "ami-15", |
| ContainerImage: "bar-arm64:latest", |
| SSHUsername: "foo", |
| }, |
| }, |
| } |
| _, err := bp.GetBuildlet(context.Background(), host, noopEventTimeLogger{}) |
| if err != nil { |
| t.Errorf("EC2Buildlet.GetBuildlet(ctx, %q, %+v) = _, %s; want no error", host, noopEventTimeLogger{}, err) |
| } |
| } |
| |
| func TestEC2BuildletGetBuildletError(t *testing.T) { |
| host := "host-type-x" |
| testCases := []struct { |
| desc string |
| hostType string |
| logger Logger |
| ledger *ledger |
| buildletClient ec2BuildletClient |
| hosts map[string]*dashboard.HostConfig |
| }{ |
| { |
| desc: "invalid-host-type", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| "wrong-host-type": {}, |
| }, |
| logger: noopEventTimeLogger{}, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: true, |
| VMCreated: true, |
| }, |
| }, |
| { |
| desc: "buildlet-client-failed-instance-created", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: {}, |
| }, |
| logger: noopEventTimeLogger{}, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: false, |
| VMCreated: false, |
| }, |
| }, |
| { |
| desc: "buildlet-client-failed-instance-not-created", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: {}, |
| }, |
| logger: noopEventTimeLogger{}, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: true, |
| VMCreated: false, |
| }, |
| }, |
| } |
| for _, tt := range testCases { |
| t.Run(tt.desc, func(t *testing.T) { |
| bp := &EC2Buildlet{ |
| buildletClient: tt.buildletClient, |
| buildEnv: &buildenv.Environment{}, |
| ledger: tt.ledger, |
| hosts: tt.hosts, |
| } |
| _, gotErr := bp.GetBuildlet(context.Background(), tt.hostType, tt.logger) |
| if gotErr == nil { |
| t.Errorf("EC2Buildlet.GetBuildlet(ctx, %q, %+v) = _, %s", tt.hostType, tt.logger, gotErr) |
| } |
| }) |
| } |
| } |
| |
| func TestEC2BuildletGetBuildletLogger(t *testing.T) { |
| host := "host-type-x" |
| testCases := []struct { |
| desc string |
| buildletClient ec2BuildletClient |
| hostType string |
| hosts map[string]*dashboard.HostConfig |
| ledger *ledger |
| wantLogs []string |
| wantSpans []string |
| wantSpansErr []string |
| }{ |
| { |
| desc: "buildlet-client-failed-instance-create-request-failed", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: {}, |
| }, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: false, |
| VMCreated: false, |
| buildletCreated: false, |
| }, |
| wantSpans: []string{"create_ec2_instance", "awaiting_ec2_quota", "create_ec2_buildlet"}, |
| wantSpansErr: []string{"create_ec2_buildlet", "create_ec2_instance"}, |
| }, |
| { |
| desc: "buildlet-client-failed-instance-not-created", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: {}, |
| }, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: true, |
| VMCreated: false, |
| buildletCreated: false, |
| }, |
| wantSpans: []string{"create_ec2_instance", "awaiting_ec2_quota", "create_ec2_buildlet"}, |
| wantSpansErr: []string{"create_ec2_buildlet", "create_ec2_instance"}, |
| }, |
| { |
| desc: "buildlet-client-failed-instance-created", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: {}, |
| }, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: true, |
| VMCreated: true, |
| buildletCreated: false, |
| }, |
| wantSpans: []string{"create_ec2_instance", "awaiting_ec2_quota", "create_ec2_buildlet", "wait_buildlet_start"}, |
| wantSpansErr: []string{"create_ec2_buildlet", "wait_buildlet_start"}, |
| }, |
| { |
| desc: "success", |
| hostType: host, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 0, |
| entries: make(map[string]*entry), |
| types: map[string]*cloud.InstanceType{ |
| "e2-highcpu-2": { |
| Type: "e2-highcpu-2", |
| CPU: 4, |
| }, |
| }, |
| }, |
| hosts: map[string]*dashboard.HostConfig{ |
| host: {}, |
| }, |
| buildletClient: &fakeEC2BuildletClient{ |
| createVMRequestSuccess: true, |
| VMCreated: true, |
| buildletCreated: true, |
| }, |
| wantSpans: []string{"create_ec2_instance", "create_ec2_buildlet", "awaiting_ec2_quota", "wait_buildlet_start"}, |
| wantSpansErr: []string{}, |
| }, |
| } |
| for _, tc := range testCases { |
| t.Run(tc.desc, func(t *testing.T) { |
| bp := &EC2Buildlet{ |
| buildletClient: tc.buildletClient, |
| buildEnv: &buildenv.Environment{}, |
| ledger: tc.ledger, |
| hosts: tc.hosts, |
| } |
| l := newTestLogger() |
| _, _ = bp.GetBuildlet(context.Background(), tc.hostType, l) |
| if !cmp.Equal(l.spanEvents(), tc.wantSpans, cmp.Transformer("sort", func(in []string) []string { |
| out := append([]string(nil), in...) |
| sort.Strings(out) |
| return out |
| })) { |
| t.Errorf("span events = %+v; want %+v", l.spanEvents(), tc.wantSpans) |
| } |
| for _, spanErr := range tc.wantSpansErr { |
| s, ok := l.spans[spanErr] |
| if !ok { |
| t.Fatalf("log span %q does not exist", spanErr) |
| } |
| if s.err == nil { |
| t.Fatalf("testLogger.span[%q].err is nil", spanErr) |
| } |
| } |
| }) |
| } |
| } |
| |
| func TestEC2BuildletString(t *testing.T) { |
| testCases := []struct { |
| desc string |
| instCount int64 |
| cpuCount int64 |
| cpuLimit int64 |
| }{ |
| {"default", 0, 0, 0}, |
| {"non-default", 2, 2, 3}, |
| } |
| for _, tc := range testCases { |
| t.Run(tc.desc, func(t *testing.T) { |
| es := make([]*entry, tc.instCount) |
| entries := make(map[string]*entry) |
| for i, e := range es { |
| entries[fmt.Sprintf("%d", i)] = e |
| } |
| eb := &EC2Buildlet{ |
| ledger: &ledger{ |
| cpuLimit: tc.cpuLimit, |
| cpuUsed: tc.cpuCount, |
| entries: entries, |
| }, |
| } |
| want := fmt.Sprintf("EC2 pool capacity: %d instances; %d/%d CPUs", tc.instCount, tc.cpuCount, tc.cpuLimit) |
| got := eb.String() |
| if got != want { |
| t.Errorf("EC2Buildlet.String() = %s; want %s", got, want) |
| } |
| }) |
| } |
| } |
| |
| func TestEC2BuildletCapacityString(t *testing.T) { |
| testCases := []struct { |
| desc string |
| instCount int64 |
| cpuCount int64 |
| cpuLimit int64 |
| }{ |
| {"defaults", 0, 0, 0}, |
| {"non-default", 2, 2, 3}, |
| } |
| for _, tc := range testCases { |
| t.Run(tc.desc, func(t *testing.T) { |
| es := make([]*entry, tc.instCount) |
| entries := make(map[string]*entry) |
| for i, e := range es { |
| entries[fmt.Sprintf("%d", i)] = e |
| } |
| eb := &EC2Buildlet{ |
| ledger: &ledger{ |
| cpuLimit: tc.cpuLimit, |
| cpuUsed: tc.cpuCount, |
| entries: entries, |
| }, |
| } |
| want := fmt.Sprintf("%d instances; %d/%d CPUs", tc.instCount, tc.cpuCount, tc.cpuLimit) |
| got := eb.capacityString() |
| if got != want { |
| t.Errorf("EC2Buildlet.capacityString() = %s; want %s", got, want) |
| } |
| }) |
| } |
| } |
| |
| func TestEC2BuildletbuildletDone(t *testing.T) { |
| t.Run("done-successful", func(t *testing.T) { |
| instName := "instance-name-x" |
| |
| awsC := cloud.NewFakeAWSClient() |
| inst, err := awsC.CreateInstance(context.Background(), &cloud.EC2VMConfiguration{ |
| Description: "test instance", |
| ImageID: "image-x", |
| Name: instName, |
| SSHKeyID: "key-14", |
| Tags: map[string]string{}, |
| Type: "type-x", |
| Zone: "zone-1", |
| }) |
| if err != nil { |
| t.Errorf("unable to create instance: %s", err) |
| } |
| |
| pool := &EC2Buildlet{ |
| awsClient: awsC, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 5, |
| entries: map[string]*entry{ |
| instName: { |
| createdAt: time.Now(), |
| instanceID: inst.ID, |
| instanceName: instName, |
| vCPUCount: 5, |
| }, |
| }, |
| }, |
| } |
| pool.buildletDone(instName) |
| if gotID := pool.ledger.InstanceID(instName); gotID != "" { |
| t.Errorf("ledger.instanceID = %q; want %q", gotID, "") |
| } |
| gotInsts, err := awsC.RunningInstances(context.Background()) |
| if err != nil || len(gotInsts) != 0 { |
| t.Errorf("awsClient.RunningInstances(ctx) = %+v, %s; want [], nil", gotInsts, err) |
| } |
| }) |
| t.Run("instance-not-in-ledger", func(t *testing.T) { |
| instName := "instance-name-x" |
| |
| awsC := cloud.NewFakeAWSClient() |
| inst, err := awsC.CreateInstance(context.Background(), &cloud.EC2VMConfiguration{ |
| Description: "test instance", |
| ImageID: "image-x", |
| Name: instName, |
| SSHKeyID: "key-14", |
| Tags: map[string]string{}, |
| Type: "type-x", |
| Zone: "zone-1", |
| }) |
| if err != nil { |
| t.Errorf("unable to create instance: %s", err) |
| } |
| |
| pool := &EC2Buildlet{ |
| awsClient: awsC, |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 5, |
| entries: map[string]*entry{}, |
| }, |
| } |
| pool.buildletDone(inst.Name) |
| gotInsts, err := awsC.RunningInstances(context.Background()) |
| if err != nil || len(gotInsts) != 1 { |
| t.Errorf("awsClient.RunningInstances(ctx) = %+v, %s; want 1 instance, nil", gotInsts, err) |
| } |
| }) |
| t.Run("instance-not-in-ec2", func(t *testing.T) { |
| instName := "instance-name-x" |
| pool := &EC2Buildlet{ |
| awsClient: cloud.NewFakeAWSClient(), |
| ledger: &ledger{ |
| cpuLimit: 20, |
| cpuUsed: 5, |
| entries: map[string]*entry{ |
| instName: { |
| createdAt: time.Now(), |
| instanceID: "instance-id-14", |
| instanceName: instName, |
| vCPUCount: 5, |
| }, |
| }, |
| }, |
| } |
| pool.buildletDone(instName) |
| if gotID := pool.ledger.InstanceID(instName); gotID != "" { |
| t.Errorf("ledger.instanceID = %q; want %q", gotID, "") |
| } |
| }) |
| } |
| |
| func TestEC2BuildletClose(t *testing.T) { |
| cancelled := false |
| pool := &EC2Buildlet{ |
| cancelPoll: func() { cancelled = true }, |
| } |
| pool.Close() |
| if !cancelled { |
| t.Error("EC2Buildlet.pollCancel not called") |
| } |
| } |
| |
| func TestEC2BuildletRetrieveAndSetQuota(t *testing.T) { |
| pool := &EC2Buildlet{ |
| awsClient: cloud.NewFakeAWSClient(), |
| ledger: newLedger(), |
| } |
| err := pool.retrieveAndSetQuota(context.Background()) |
| if err != nil { |
| t.Errorf("EC2Buildlet.retrieveAndSetQuota(ctx) = %s; want nil", err) |
| } |
| if pool.ledger.cpuLimit == 0 { |
| t.Errorf("ledger.cpuLimit = %d; want non-zero", pool.ledger.cpuLimit) |
| } |
| } |
| |
| func TestEC2BuildletRetrieveAndSetInstanceTypes(t *testing.T) { |
| pool := &EC2Buildlet{ |
| awsClient: cloud.NewFakeAWSClient(), |
| ledger: newLedger(), |
| } |
| err := pool.retrieveAndSetInstanceTypes() |
| if err != nil { |
| t.Errorf("EC2Buildlet.retrieveAndSetInstanceTypes() = %s; want nil", err) |
| } |
| if len(pool.ledger.types) == 0 { |
| t.Errorf("len(pool.ledger.types) = %d; want non-zero", len(pool.ledger.types)) |
| } |
| } |
| |
| func TestEC2BuildeletDestroyUntrackedInstances(t *testing.T) { |
| awsC := cloud.NewFakeAWSClient() |
| create := func(name string) *cloud.Instance { |
| inst, err := awsC.CreateInstance(context.Background(), &cloud.EC2VMConfiguration{ |
| Description: "test instance", |
| ImageID: "image-x", |
| Name: name, |
| SSHKeyID: "key-14", |
| Tags: map[string]string{}, |
| Type: "type-x", |
| Zone: "zone-1", |
| }) |
| if err != nil { |
| t.Errorf("unable to create instance: %s", err) |
| } |
| return inst |
| } |
| // create untracked instances |
| for it := 0; it < 10; it++ { |
| _ = create(instanceName("host-test-type", 10)) |
| } |
| wantTrackedInst := create(instanceName("host-test-type", 10)) |
| wantRemoteInst := create(instanceName("host-test-type", 10)) |
| _ = create("debug-tiger-host-14") // non buildlet instance |
| |
| pool := &EC2Buildlet{ |
| awsClient: awsC, |
| isRemoteBuildlet: func(name string) bool { |
| if name == wantRemoteInst.Name { |
| return true |
| } |
| return false |
| }, |
| ledger: &ledger{ |
| cpuLimit: 200, |
| cpuUsed: 4, |
| entries: map[string]*entry{ |
| wantTrackedInst.Name: { |
| createdAt: time.Now(), |
| instanceID: wantTrackedInst.ID, |
| instanceName: wantTrackedInst.Name, |
| vCPUCount: 4, |
| }, |
| }, |
| }, |
| } |
| pool.destroyUntrackedInstances(context.Background()) |
| wantInstCount := 3 |
| gotInsts, err := awsC.RunningInstances(context.Background()) |
| if err != nil || len(gotInsts) != wantInstCount { |
| t.Errorf("awsClient.RunningInstances(ctx) = %+v, %s; want %d instances and no error", gotInsts, err, wantInstCount) |
| } |
| } |
| |
| // fakeEC2BuildletClient is the client used to create buildlets on EC2. |
| type fakeEC2BuildletClient struct { |
| createVMRequestSuccess bool |
| VMCreated bool |
| buildletCreated bool |
| } |
| |
| // StartNewVM boots a new VM on EC2, waits until the client is accepting connections |
| // on the configured port and returns a buildlet client configured communicate with it. |
| func (f *fakeEC2BuildletClient) StartNewVM(ctx context.Context, buildEnv *buildenv.Environment, hconf *dashboard.HostConfig, vmName, hostType string, opts *buildlet.VMOpts) (buildlet.Client, error) { |
| // check required params |
| if opts == nil || opts.TLS.IsZero() { |
| return nil, errors.New("TLS keypair is not set") |
| } |
| if buildEnv == nil { |
| return nil, errors.New("invalid build enviornment") |
| } |
| if hconf == nil { |
| return nil, errors.New("invalid host configuration") |
| } |
| if vmName == "" || hostType == "" { |
| return nil, fmt.Errorf("invalid vmName: %q and hostType: %q", vmName, hostType) |
| } |
| if opts.DeleteIn == 0 { |
| // Note: This implements a short default in the rare case the caller doesn't care. |
| opts.DeleteIn = 30 * time.Minute |
| } |
| if !f.createVMRequestSuccess { |
| return nil, fmt.Errorf("unable to create instance %s: creation disabled", vmName) |
| } |
| condRun := func(fn func()) { |
| if fn != nil { |
| fn() |
| } |
| } |
| condRun(opts.OnInstanceRequested) |
| if !f.VMCreated { |
| return nil, errors.New("error waiting for instance to exist: vm existance disabled") |
| } |
| |
| condRun(opts.OnInstanceCreated) |
| |
| if !f.buildletCreated { |
| return nil, errors.New("error waiting for buildlet: buildlet creation disabled") |
| } |
| |
| if opts.OnGotEC2InstanceInfo != nil { |
| opts.OnGotEC2InstanceInfo(&cloud.Instance{ |
| CPUCount: 4, |
| CreatedAt: time.Time{}, |
| Description: "sample vm", |
| ID: "id-" + instanceName("random", 4), |
| IPAddressExternal: "127.0.0.1", |
| IPAddressInternal: "127.0.0.1", |
| ImageID: "image-x", |
| Name: vmName, |
| SSHKeyID: "key-15", |
| SecurityGroups: nil, |
| State: "running", |
| Tags: map[string]string{ |
| "foo": "bar", |
| }, |
| Type: "yy.large", |
| Zone: "zone-a", |
| }) |
| } |
| return &buildlet.FakeClient{}, nil |
| } |
| |
| type testLogger struct { |
| eventTimes []eventTime |
| spans map[string]*span |
| } |
| |
| type eventTime struct { |
| event string |
| opt []string |
| } |
| |
| type span struct { |
| event string |
| opt []string |
| err error |
| calledDone bool |
| } |
| |
| func (s *span) Done(err error) error { |
| s.err = err |
| s.calledDone = true |
| return nil |
| } |
| |
| func newTestLogger() *testLogger { |
| return &testLogger{ |
| eventTimes: make([]eventTime, 0, 5), |
| spans: make(map[string]*span), |
| } |
| } |
| |
| func (l *testLogger) LogEventTime(event string, optText ...string) { |
| l.eventTimes = append(l.eventTimes, eventTime{ |
| event: event, |
| opt: optText, |
| }) |
| } |
| |
| func (l *testLogger) CreateSpan(event string, optText ...string) spanlog.Span { |
| s := &span{ |
| event: event, |
| opt: optText, |
| } |
| l.spans[event] = s |
| return s |
| } |
| |
| func (l *testLogger) spanEvents() []string { |
| se := make([]string, 0, len(l.spans)) |
| for k, s := range l.spans { |
| if !s.calledDone { |
| continue |
| } |
| se = append(se, k) |
| } |
| return se |
| } |
| |
| type noopEventTimeLogger struct{} |
| |
| func (l noopEventTimeLogger) LogEventTime(event string, optText ...string) {} |
| func (l noopEventTimeLogger) CreateSpan(event string, optText ...string) spanlog.Span { |
| return noopSpan{} |
| } |
| |
| type noopSpan struct{} |
| |
| func (s noopSpan) Done(err error) error { return nil } |