diff --git a/pkg/provider/aws/aws.go b/pkg/provider/aws/aws.go index 7838536..ef83b61 100644 --- a/pkg/provider/aws/aws.go +++ b/pkg/provider/aws/aws.go @@ -123,6 +123,9 @@ type Client interface { // ListInstances returns every Nebula-owned instance (filtered by the // ClaimTagKey tag) across the region, in as few calls as possible. ListInstances(ctx context.Context) ([]EC2Instance, error) + // FindInstance returns the live instance tagged with claimName, or (nil, nil) if none. + // Filtered server-side and without status checks, since only the id is read. + FindInstance(ctx context.Context, claimName string) (*EC2Instance, error) // AvailableInstanceTypes returns the set of EC2 instance types the client's // region actually offers, as a set keyed by instance type. It backs the // per-region availability filter in Offerings: a static catalog row whose @@ -719,16 +722,7 @@ func (p *Provider) ClassifyProvisionError(err error, accelerator, region string) // or nil if none. It scans one region's client (the launch target), since a claim // is placed in exactly one region per provision attempt. func findByClaim(ctx context.Context, client Client, claimName string) (*EC2Instance, error) { - instances, err := client.ListInstances(ctx) - if err != nil { - return nil, err - } - for i := range instances { - if instances[i].Tags[ClaimTagKey] == claimName { - return &instances[i], nil - } - } - return nil, nil + return client.FindInstance(ctx, claimName) } // instanceSpecFromPod reads the workload off the Pod (source of truth) and the diff --git a/pkg/provider/aws/aws_test.go b/pkg/provider/aws/aws_test.go index 1c180ef..4921d33 100644 --- a/pkg/provider/aws/aws_test.go +++ b/pkg/provider/aws/aws_test.go @@ -45,6 +45,7 @@ type fakeClient struct { instances []EC2Instance lastSpec InstanceSpec runCnt int + listCnt int runErr error runID string terminated []string @@ -91,9 +92,20 @@ func (f *fakeClient) DescribeInstance(_ context.Context, id string) (*EC2Instanc } func (f *fakeClient) ListInstances(_ context.Context) ([]EC2Instance, error) { + f.listCnt++ return f.instances, nil } +func (f *fakeClient) FindInstance(_ context.Context, claimName string) (*EC2Instance, error) { + for i := range f.instances { + if f.instances[i].Tags[ClaimTagKey] == claimName { + inst := f.instances[i] + return &inst, nil + } + } + return nil, nil +} + func (f *fakeClient) AvailableInstanceTypes(_ context.Context) (map[string]bool, error) { if f.availErr != nil { return nil, f.availErr @@ -390,6 +402,10 @@ func TestProvision_Idempotent(t *testing.T) { if f.runCnt != 0 { t.Fatalf("RunInstance called %d times, want 0 (idempotent)", f.runCnt) } + // A full list returns every Nebula instance in the region plus their status checks. + if f.listCnt != 0 { + t.Errorf("ListInstances called %d times on Provision, want 0", f.listCnt) + } } func TestProvision_UnsupportedAccelerator(t *testing.T) { diff --git a/pkg/provider/aws/client.go b/pkg/provider/aws/client.go index abdd618..5022b66 100644 --- a/pkg/provider/aws/client.go +++ b/pkg/provider/aws/client.go @@ -727,6 +727,40 @@ func (c *sdkClient) DescribeInstance(ctx context.Context, id string) (*EC2Instan return nil, nil } +// FindInstance implements Client. The state filter matters as much as the tag: claim names +// are reused across Pod restarts, and a terminated instance stays visible for ~1h. Unlike +// ListInstances it drops stopped instances too: toState reads them as Terminated, so +// adopting one would fail the Pod instead of launching a replacement. +func (c *sdkClient) FindInstance(ctx context.Context, claimName string) (*EC2Instance, error) { + out, err := c.ec2.DescribeInstances(ctx, &ec2.DescribeInstancesInput{ + Filters: []ec2types.Filter{ + { + Name: awssdk.String("tag:" + ClaimTagKey), + Values: []string{claimName}, + }, + stateFilter(statePending, stateRunning), + }, + }) + if err != nil { + return nil, err + } + for _, r := range out.Reservations { + if len(r.Instances) > 0 { + inst := c.observe(r.Instances[0]) + return &inst, nil + } + } + return nil, nil +} + +// stateFilter restricts DescribeInstances to the given instance states. +func stateFilter(states ...string) ec2types.Filter { + return ec2types.Filter{ + Name: awssdk.String("instance-state-name"), + Values: states, + } +} + // ListInstances implements Client. It scopes the list server-side to instances // carrying the ClaimTagKey tag — the tag every Nebula instance is launched with — // so instances Nebula does not own are never returned, and pages through all @@ -745,10 +779,7 @@ func (c *sdkClient) ListInstances(ctx context.Context) ([]EC2Instance, error) { Name: awssdk.String("tag-key"), Values: []string{ClaimTagKey}, }, - { - Name: awssdk.String("instance-state-name"), - Values: []string{"pending", "running", "stopping", "stopped"}, - }, + stateFilter(statePending, stateRunning, stateStopping, stateStopped), }, } var out []EC2Instance diff --git a/pkg/provider/aws/client_test.go b/pkg/provider/aws/client_test.go index 041ec66..fbca1e6 100644 --- a/pkg/provider/aws/client_test.go +++ b/pkg/provider/aws/client_test.go @@ -1089,6 +1089,45 @@ func TestSDKAvailableInstanceTypes_ErrorPropagates(t *testing.T) { } } +func TestSDKFindInstance_FiltersByClaimValue(t *testing.T) { + f := &fakeEC2{describePages: []*ec2.DescribeInstancesOutput{{ + Reservations: []ec2types.Reservation{{Instances: []ec2types.Instance{{ + InstanceId: awssdk.String("i-1"), + State: &ec2types.InstanceState{Name: ec2types.InstanceStateNameRunning}, + }}}}, + }}} + c := newSDKClient(f) + + inst, err := c.FindInstance(context.Background(), "claim-a") + if err != nil { + t.Fatalf("FindInstance: %v", err) + } + if inst == nil || inst.ID != "i-1" { + t.Fatalf("inst = %+v, want i-1", inst) + } + filters := map[string][]string{} + for _, fl := range f.lastDescribe.Filters { + filters[awssdk.ToString(fl.Name)] = fl.Values + } + if got := filters["tag:"+ClaimTagKey]; len(got) != 1 || got[0] != "claim-a" { + t.Errorf("claim filter = %v, want [claim-a]", got) + } + // Claim names are reused across Pod restarts, so anything else lets a retry adopt a + // terminated instance EC2 still shows, or a stopped one toState reads as Terminated. + if got := filters["instance-state-name"]; !reflect.DeepEqual(got, []string{statePending, stateRunning}) { + t.Errorf("instance-state-name filter = %v, want [pending running]", got) + } + // Provision reads only the id, so the status probe would be a wasted call. + if f.lastStatusIn != nil { + t.Error("DescribeInstanceStatus called; FindInstance must not probe status checks") + } + + none, err := newSDKClient(&fakeEC2{}).FindInstance(context.Background(), "claim-b") + if err != nil || none != nil { + t.Errorf("FindInstance(no match) = %v, %v; want nil, nil", none, err) + } +} + func TestSDKList_PagesAndFiltersByClaimTag(t *testing.T) { f := &fakeEC2{describePages: []*ec2.DescribeInstancesOutput{ { diff --git a/pkg/provider/modal/client.go b/pkg/provider/modal/client.go index d43944a..f1a5c8f 100644 --- a/pkg/provider/modal/client.go +++ b/pkg/provider/modal/client.go @@ -716,6 +716,33 @@ func (c *sdkClient) ListSandboxes(ctx context.Context) ([]Sandbox, error) { return out, nil } +// FindSandbox implements Client. Modal's Tags filter matches exact key=value pairs, which is +// useless for ListSandboxes (every claim value differs) but exactly one claim's lookup. +func (c *sdkClient) FindSandbox(ctx context.Context, claimName string) (*Sandbox, error) { + app, err := c.app(ctx) + if err != nil { + return nil, fmt.Errorf("modal: resolve app: %w", err) + } + seq, err := c.mc.Sandboxes.List(ctx, &modal.SandboxListParams{ + AppID: app.AppID, + Tags: map[string]string{ClaimTagKey: claimName}, + }) + if err != nil { + return nil, err + } + for sb, err := range seq { + if err != nil { + return nil, err + } + observed, err := c.observe(ctx, sb) + if err != nil { + return nil, err + } + return &observed, nil + } + return nil, nil +} + // observe normalizes a live SDK *Sandbox into the adapter-level Sandbox view: tags // (from GetTags) and status (from Poll). A Poll error is tolerated so a single flaky // sandbox doesn't fail the whole read — the poll loop will re-observe next tick. A TAG diff --git a/pkg/provider/modal/modal.go b/pkg/provider/modal/modal.go index e945fb9..85dbf2e 100644 --- a/pkg/provider/modal/modal.go +++ b/pkg/provider/modal/modal.go @@ -119,6 +119,9 @@ type Client interface { // ListSandboxes returns every Nebula-owned sandbox, filtered by the tag the // adapter sets at create time, in as few calls as possible. ListSandboxes(ctx context.Context) ([]Sandbox, error) + // FindSandbox returns the live sandbox tagged with claimName, or (nil, nil) if none. + // Filtered server-side, so it costs one list call and one observe, not one per sandbox. + FindSandbox(ctx context.Context, claimName string) (*Sandbox, error) // SandboxLogs returns merged stdout+stderr, from the first byte, following until // the sandbox exits (see provider.LogStreamer). The caller owns Close. SandboxLogs(ctx context.Context, id string) (io.ReadCloser, error) @@ -593,18 +596,12 @@ func (p *Provider) ClassifyProvisionError(err error, accelerator, region string) // findByClaim returns the sandbox tagged with claimName, or nil if none. func (p *Provider) findByClaim(ctx context.Context, claimName string) (*provider.Instance, error) { - // TODO: do we have performance issue here? - sandboxes, err := p.client.ListSandboxes(ctx) - if err != nil { + sb, err := p.client.FindSandbox(ctx, claimName) + if err != nil || sb == nil { return nil, err } - for _, sb := range sandboxes { - if sb.Tags[ClaimTagKey] == claimName { - inst := p.toInstance(sb) - return &inst, nil - } - } - return nil, nil + inst := p.toInstance(*sb) + return &inst, nil } // sandboxSpecFromPod reads the workload off the Pod (source of truth) and the diff --git a/pkg/provider/modal/modal_test.go b/pkg/provider/modal/modal_test.go index f211871..c15a5e0 100644 --- a/pkg/provider/modal/modal_test.go +++ b/pkg/provider/modal/modal_test.go @@ -45,6 +45,7 @@ type fakeClient struct { sandboxes []Sandbox lastSpec SandboxSpec createCnt int + listCnt int createErr error createID string cred Credential // credential CreateSandbox returns alongside its id @@ -111,9 +112,20 @@ func (f *fakeClient) GetSandbox(_ context.Context, id string) (*Sandbox, error) } func (f *fakeClient) ListSandboxes(_ context.Context) ([]Sandbox, error) { + f.listCnt++ return f.sandboxes, nil } +func (f *fakeClient) FindSandbox(_ context.Context, claimName string) (*Sandbox, error) { + for i := range f.sandboxes { + if f.sandboxes[i].Tags[ClaimTagKey] == claimName { + s := f.sandboxes[i] + return &s, nil + } + } + return nil, nil +} + func (f *fakeClient) SandboxLogs(_ context.Context, id string) (io.ReadCloser, error) { f.logsFor = id if f.logsErr != nil { @@ -458,6 +470,11 @@ func TestProvision_Idempotent(t *testing.T) { if f.createCnt != 0 { t.Fatalf("CreateSandbox called %d times, want 0 (idempotent)", f.createCnt) } + // A full list observes every sandbox in the App (two calls each), so the lookup must + // go through the claim-filtered FindSandbox instead. + if f.listCnt != 0 { + t.Errorf("ListSandboxes called %d times on Provision, want 0", f.listCnt) + } // An adopted sandbox has been OBSERVED, unlike a fresh create, so its state is // known: this one is running, which means capacity was necessarily allocated. if !reserved {