Remove Slurm access the allocations no longer grant

Fixes #540

Author: lahirujayathilake

connectors/SLURM/Association-Mapper/internal/subscribers/association_writer.go

@@ -74,71 +74,88 @@ func (a *AssociationSubscriber) upsertAssociationsForMembership(ctx context.Cont
 	return a.syncAssociationsForMembership(ctx, membership, nil)
 }
 
-// syncAssociationsForMembership builds the association for each of the
-// allocation's resources and writes the ones that differ from existing, the
-// state already on the cluster. A nil existing writes all of them.
+// syncAssociationsForMembership writes the associations this membership should
+// have, skipping any that already match existing, the state already on the
+// cluster. Pass nil to skip that comparison and write every record.
 //
 // Every caller goes through here on purpose. Slurm keys an association by
 // (cluster, account, user, partition) and its upsert is last-write-wins, so a
 // caller that built the record without limits would wipe limits another caller
 // had set. Folding the per-member overrides in here keeps all writers
 // producing the same record for the same key.
 func (a *AssociationSubscriber) syncAssociationsForMembership(ctx context.Context, membership models.ComputeAllocationMembership, existing map[assocKey]client.Association) error {
+	desired, err := a.desiredAssociationsForMembership(ctx, membership)
+	if err != nil {
+		return err
+	}
+	for _, association := range desired {
+		if got, ok := existing[keyOf(association)]; ok && sameLimits(association, got) {
+			continue
+		}
+		if err := a.slurmClient.UpsertAssociation(association); err != nil {
+			return fmt.Errorf("upsert association for partition %s: %w", association.Partition, err)
+		}
+		slog.Info("Upserted association", "association", association)
+	}
+	return nil
+}
+
+// desiredAssociationsForMembership builds the associations this membership
+// should have on the cluster, without writing anything. An inactive allocation
+// grants nothing, so it yields an empty set rather than an error.
+func (a *AssociationSubscriber) desiredAssociationsForMembership(ctx context.Context, membership models.ComputeAllocationMembership) ([]client.Association, error) {
 	allocation, err := a.coreService.GetComputeAllocation(ctx, membership.ComputeAllocationID)
 	if err != nil {
-		return fmt.Errorf("get compute allocation: %w", err)
+		return nil, fmt.Errorf("get compute allocation: %w", err)
+	}
+	if !activeAllocation(*allocation) {
+		return nil, nil
 	}
 	cluster, err := a.coreService.GetComputeCluster(ctx, allocation.ComputeClusterID)
 	if err != nil {
-		return fmt.Errorf("get compute cluster: %w", err)
+		return nil, fmt.Errorf("get compute cluster: %w", err)
 	}
 	csu, err := a.coreService.GetComputeClusterUserByPair(ctx, cluster.ID, membership.UserID)
 	if err != nil {
-		return fmt.Errorf("get compute cluster user: %w", err)
+		return nil, fmt.Errorf("get compute cluster user: %w", err)
 	}
 	if csu.ProvisionedAt == nil {
-		return errNotProvisioned
+		return nil, errNotProvisioned
 	}
 
 	resources, err := a.coreService.ListResourcesForAllocation(ctx, allocation.ID)
 	if err != nil {
-		return fmt.Errorf("list resources for allocation: %w", err)
+		return nil, fmt.Errorf("list resources for allocation: %w", err)
 	}
 	if len(resources) == 0 {
 		// Nothing to map onto a partition. Skipping beats guessing a
 		// partition name the cluster may not have.
 		slog.Warn("Allocation has no resources, no association written",
 			"allocation_id", allocation.ID, "user_id", membership.UserID)
-		return nil
+		return nil, nil
 	}
 
 	overrides, err := a.coreService.ListOverridesForMembership(ctx, membership.ID)
 	if err != nil {
-		return fmt.Errorf("list overrides for membership: %w", err)
+		return nil, fmt.Errorf("list overrides for membership: %w", err)
 	}
 	overrideByResource := make(map[string]models.ComputeAllocationMembershipResourceOverride, len(overrides))
 	for _, o := range overrides {
 		overrideByResource[o.ComputeAllocationResourceID] = o
 	}
 
+	out := make([]client.Association, 0, len(resources))
 	for _, resource := range resources {
-		association := client.Association{
+		out = append(out, client.Association{
 			Account:   allocation.Name,
 			Cluster:   cluster.Name,
 			User:      csu.LocalUsername,
 			Partition: resource.Name,
 			QoS:       []string{"normal"},
 			Limits:    limitsFor(resource, overrideByResource[resource.ID]),
-		}
-		if got, ok := existing[keyOf(association)]; ok && sameLimits(association, got) {
-			continue
-		}
-		if err := a.slurmClient.UpsertAssociation(association); err != nil {
-			return fmt.Errorf("upsert association for partition %s: %w", resource.Name, err)
-		}
-		slog.Info("Upserted association", "association", association)
+		})
 	}
-	return nil
+	return out, nil
 }
 
 // limitsFor turns a per-member override into association limits. A zero

connectors/SLURM/Association-Mapper/internal/subscribers/deprovision_test.go

@@ -22,6 +22,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/apache/airavata-custos/connectors/SLURM/Rest-Client/pkg/client"
 	"github.com/apache/airavata-custos/pkg/models"
 )
 
@@ -153,3 +154,116 @@ func TestAllocationDeletionRemovesAllAssociations(t *testing.T) {
 		t.Fatalf("expected an account-wide delete on allocation deletion, got %+v", got)
 	}
 }
+
+// The sweep is the backstop for a lost deactivation: an association the
+// allocations no longer call for is removed even if no event ever arrived.
+func TestReconcilerRemovesStaleAssociation(t *testing.T) {
+	core := coreMock(mockOpts{
+		provisionedAt: ago(time.Minute),
+		memberships:   []models.ComputeAllocationMembership{testMembership()},
+	})
+	slurm := &fakeSlurmClient{existing: []client.Association{
+		// desired
+		{Account: "test-alloc", Cluster: "testcluster", User: "testuser", Partition: "compute"},
+		// left behind by a membership that is gone
+		{Account: "test-alloc", Cluster: "testcluster", User: "ghost", Partition: "compute"},
+	}}
+	NewAssociationSubscriber(slurm, nil, core, 0, 0).reconcile(context.Background())
+
+	got := slurm.allDeletes()
+	if len(got) != 1 {
+		t.Fatalf("expected the stale association to be removed, got %d deletes", len(got))
+	}
+	if got[0].User != "ghost" || got[0].Partition != "compute" {
+		t.Errorf("removed the wrong association: %+v", got[0])
+	}
+}
+
+// Guard: an empty desired set almost always means a failed lookup, so the
+// sweep must not read it as "revoke everyone".
+func TestReconcilerDoesNotPruneWhenNothingIsDesired(t *testing.T) {
+	core := coreMock(mockOpts{
+		provisionedAt: ago(time.Minute),
+		memberships:   nil, // nobody entitled to anything
+	})
+	slurm := &fakeSlurmClient{existing: []client.Association{
+		{Account: "test-alloc", Cluster: "testcluster", User: "testuser", Partition: "compute"},
+	}}
+	NewAssociationSubscriber(slurm, nil, core, 0, 0).reconcile(context.Background())
+
+	if n := len(slurm.allDeletes()); n != 0 {
+		t.Fatalf("an empty desired set must not revoke anything, got %d deletes", n)
+	}
+}
+
+// Guard: associations on accounts Custos does not manage are never touched.
+func TestReconcilerLeavesUnmanagedAccountsAlone(t *testing.T) {
+	core := coreMock(mockOpts{
+		provisionedAt:     ago(time.Minute),
+		memberships:       []models.ComputeAllocationMembership{testMembership()},
+		unmanagedAccounts: true, // core reports no allocations on this cluster
+	})
+	slurm := &fakeSlurmClient{existing: []client.Association{
+		{Account: "someone-elses-account", Cluster: "testcluster", User: "outsider", Partition: "compute"},
+	}}
+	NewAssociationSubscriber(slurm, nil, core, 0, 0).reconcile(context.Background())
+
+	if n := len(slurm.allDeletes()); n != 0 {
+		t.Fatalf("an unmanaged account must not be touched, got %d deletes", n)
+	}
+}
+
+// Account-level records carry the allocation's own limits, not a member's
+// access, so the sweep must leave them be.
+func TestReconcilerLeavesAccountLevelAssociationsAlone(t *testing.T) {
+	core := coreMock(mockOpts{
+		provisionedAt: ago(time.Minute),
+		memberships:   []models.ComputeAllocationMembership{testMembership()},
+	})
+	slurm := &fakeSlurmClient{existing: []client.Association{
+		{Account: "test-alloc", Cluster: "testcluster", User: "testuser", Partition: "compute"},
+		{Account: "test-alloc", Cluster: "testcluster", User: ""}, // account-level
+	}}
+	NewAssociationSubscriber(slurm, nil, core, 0, 0).reconcile(context.Background())
+
+	if n := len(slurm.allDeletes()); n != 0 {
+		t.Fatalf("account-level associations must not be pruned, got %d deletes", n)
+	}
+}
+
+// A member inside the provisioning grace is still entitled, so their fresh
+// association must not be pruned just because the sweep is not writing it yet.
+func TestReconcilerDoesNotPruneAssociationsInsideGrace(t *testing.T) {
+	core := coreMock(mockOpts{
+		provisionedAt: ago(time.Second),
+		memberships:   []models.ComputeAllocationMembership{testMembership()},
+	})
+	slurm := &fakeSlurmClient{existing: []client.Association{
+		{Account: "test-alloc", Cluster: "testcluster", User: "testuser", Partition: "compute"},
+	}}
+	NewAssociationSubscriber(slurm, nil, core, 0, 0).reconcile(context.Background())
+
+	if n := len(slurm.allDeletes()); n != 0 {
+		t.Fatalf("a member inside the grace must keep their association, got %d deletes", n)
+	}
+}
+
+// An allocation that is no longer active grants nothing, so its members'
+// associations are swept away even without a deactivation event.
+func TestReconcilerRemovesAssociationsForInactiveAllocation(t *testing.T) {
+	core := coreMock(mockOpts{
+		provisionedAt:    ago(time.Minute),
+		allocationStatus: models.INACTIVE,
+		memberships:      []models.ComputeAllocationMembership{testMembership()},
+	})
+	slurm := &fakeSlurmClient{existing: []client.Association{
+		{Account: "test-alloc", Cluster: "testcluster", User: "testuser", Partition: "compute"},
+	}}
+	NewAssociationSubscriber(slurm, nil, core, 0, 0).reconcile(context.Background())
+
+	// Desired is empty for an inactive allocation, so the empty-desired guard
+	// holds and nothing is revoked. The event handler does that job.
+	if n := len(slurm.allDeletes()); n != 0 {
+		t.Fatalf("expected the empty-desired guard to hold, got %d deletes", n)
+	}
+}

connectors/SLURM/Association-Mapper/internal/subscribers/members_test.go

@@ -69,20 +69,26 @@ func (f *fakeSlurmClient) all() []client.Association {
 }
 
 type mockOpts struct {
-	provisionedAt *time.Time
-	resources     []models.ComputeAllocationResource
-	overrides     []models.ComputeAllocationMembershipResourceOverride
-	memberships   []models.ComputeAllocationMembership
+	allocationStatus  models.AllocationStatus
+	unmanagedAccounts bool
+	provisionedAt     *time.Time
+	resources         []models.ComputeAllocationResource
+	overrides         []models.ComputeAllocationMembershipResourceOverride
+	memberships       []models.ComputeAllocationMembership
 }
 
 func coreMock(o mockOpts) *service.CoreServiceMock {
+	allocStatus := o.allocationStatus
+	if allocStatus == "" {
+		allocStatus = models.ACTIVE
+	}
 	resources := o.resources
 	if resources == nil {
 		resources = []models.ComputeAllocationResource{{ID: "res-1", Name: "compute", ResourceType: "cpu"}}
 	}
 	return &service.CoreServiceMock{
 		GetComputeAllocationFunc: func(ctx context.Context, id string) (*models.ComputeAllocation, error) {
-			return &models.ComputeAllocation{ID: id, Name: "test-alloc", ComputeClusterID: "cluster-1"}, nil
+			return &models.ComputeAllocation{ID: id, Name: "test-alloc", ComputeClusterID: "cluster-1", Status: allocStatus}, nil
 		},
 		GetComputeClusterFunc: func(ctx context.Context, id string) (*models.ComputeCluster, error) {
 			return &models.ComputeCluster{ID: id, Name: "testcluster"}, nil
@@ -102,6 +108,14 @@ func coreMock(o mockOpts) *service.CoreServiceMock {
 				LocalUsername: "testuser", ProvisionedAt: o.provisionedAt,
 			}, nil
 		},
+		ListComputeAllocationsByClusterFunc: func(ctx context.Context, clusterID string) ([]models.ComputeAllocation, error) {
+			if o.unmanagedAccounts {
+				return nil, nil
+			}
+			return []models.ComputeAllocation{
+				{ID: "alloc-1", Name: "test-alloc", ComputeClusterID: clusterID, Status: allocStatus},
+			}, nil
+		},
 		ListComputeClustersFunc: func(ctx context.Context) ([]models.ComputeCluster, error) {
 			return []models.ComputeCluster{{ID: "cluster-1", Name: "testcluster"}}, nil
 		},