Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 19 additions & 11 deletions server/api/services/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -261,19 +261,23 @@ func (s *service) deviceCleanup() store.TransactionCb {
uids := make([]string, len(devices))
for i, device := range devices {
uids[i] = device.UID
totalDeleted++
deletedPerTenant[device.TenantID]++
}

if _, err := s.store.DeviceDeleteMany(ctx, uids); err != nil {
deleted, err := s.store.DeviceDeleteMany(ctx, uids)
if err != nil {
log.WithField("page", page+1).
WithError(err).
Info("Failed to delete devices")

return err
}

log.WithFields(log.Fields{"page": page + 1, "total_pages": totalPages, "devices": len(devices), "total_deleted": totalDeleted}).
for _, device := range deleted {
totalDeleted++
deletedPerTenant[device.TenantID]++
}

log.WithFields(log.Fields{"page": page + 1, "total_pages": totalPages, "listed": len(devices), "deleted": len(deleted), "total_deleted": totalDeleted}).
Info("Processing page of removed devices")

if page < totalPages-1 {
Expand Down Expand Up @@ -318,21 +322,25 @@ func (s *service) ephemeralCleanup() store.TransactionCb {
log.WithField("total_devices", len(devices)).Info("Found offline ephemeral devices to cleanup")

uids := make([]string, len(devices))
deletedPerTenant := make(map[string]map[models.DeviceStatus]int64)
for i, device := range devices {
uids[i] = device.UID
if deletedPerTenant[device.TenantID] == nil {
deletedPerTenant[device.TenantID] = make(map[models.DeviceStatus]int64)
}
deletedPerTenant[device.TenantID][device.Status]++
}

if _, err := s.store.DeviceDeleteMany(ctx, uids); err != nil {
deleted, err := s.store.DeviceDeleteMany(ctx, uids)
if err != nil {
log.WithError(err).Error("Failed to delete offline ephemeral devices")

return err
}

deletedPerTenant := make(map[string]map[models.DeviceStatus]int64)
for _, device := range deleted {
if deletedPerTenant[device.TenantID] == nil {
deletedPerTenant[device.TenantID] = make(map[models.DeviceStatus]int64)
}
deletedPerTenant[device.TenantID][device.Status]++
}

for _, tenantID := range slices.Sorted(maps.Keys(deletedPerTenant)) {
for _, status := range slices.Sorted(maps.Keys(deletedPerTenant[tenantID])) {
count := deletedPerTenant[tenantID][status]
Expand All @@ -345,7 +353,7 @@ func (s *service) ephemeralCleanup() store.TransactionCb {
}
}

log.WithField("total_deleted", len(devices)).Info("Ephemeral device cleanup completed successfully")
log.WithField("total_deleted", len(deleted)).Info("Ephemeral device cleanup completed successfully")

return nil
}
Expand Down
142 changes: 137 additions & 5 deletions server/api/services/task_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ func TestService_DeviceCleanup(t *testing.T) {
Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2"}).
Return(int64(0), errors.New("delete error")).
Return(nil, errors.New("delete error")).
Once()
},
expected: errors.New("delete error"),
Expand Down Expand Up @@ -215,7 +215,7 @@ func TestService_DeviceCleanup(t *testing.T) {
Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2", "device-3"}).
Return(int64(3), nil).
Return([]store.DeletedDevice{{UID: "device-1", TenantID: "tenant-1", Status: models.DeviceStatusRemoved}, {UID: "device-2", TenantID: "tenant-1", Status: models.DeviceStatusRemoved}, {UID: "device-3", TenantID: "tenant-2", Status: models.DeviceStatusRemoved}}, nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-1"), models.DeviceStatusRemoved, int64(-2)).
Expand Down Expand Up @@ -265,7 +265,7 @@ func TestService_DeviceCleanup(t *testing.T) {
Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2", "device-3"}).
Return(int64(3), nil).
Return([]store.DeletedDevice{{UID: "device-1", TenantID: "tenant-1", Status: models.DeviceStatusRemoved}, {UID: "device-2", TenantID: "tenant-1", Status: models.DeviceStatusRemoved}, {UID: "device-3", TenantID: "tenant-2", Status: models.DeviceStatusRemoved}}, nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-1"), models.DeviceStatusRemoved, int64(-2)).
Expand All @@ -278,6 +278,52 @@ func TestService_DeviceCleanup(t *testing.T) {
},
expected: nil,
},
{
description: "decrements the counters only by the devices it deleted",
requiredMocks: func() {
queryOptionsMock.
On("Match", mock.MatchedBy(matchFilter())).
Return(nil).
Once()
storeMock.
On("DeviceList", ctx, mock.Anything, store.DeviceAcceptableAsFalse, countOpts).
Return([]models.Device{}, int64(3), nil).
Once()
queryOptionsMock.
On("Match", mock.MatchedBy(matchFilter())).
Return(nil).
Once()
queryOptionsMock.
On("Sort", &sorter).
Return(nil).
Once()
queryOptionsMock.
On("Paginate", &query.Paginator{Page: 1, PerPage: 1000}).
Return(nil).
Once()
storeMock.
On("DeviceList", ctx, mock.Anything, store.DeviceAcceptableAsFalse, pageOpts).
Return(
[]models.Device{
{UID: "device-1", TenantID: "tenant-1", RemovedAt: &thirtyDaysAgo},
{UID: "device-2", TenantID: "tenant-1", RemovedAt: &thirtyDaysAgo},
{UID: "device-3", TenantID: "tenant-2", RemovedAt: &thirtyDaysAgo},
},
int64(3),
nil,
).
Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2", "device-3"}).
Return([]store.DeletedDevice{{UID: "device-2", TenantID: "tenant-1", Status: models.DeviceStatusRemoved}}, nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-1"), models.DeviceStatusRemoved, int64(-1)).
Return(nil).
Once()
},
expected: nil,
},
{
description: "succeeds with multiple pages and updates counters correctly",
requiredMocks: func() {
Expand Down Expand Up @@ -313,7 +359,7 @@ func TestService_DeviceCleanup(t *testing.T) {
Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1"}).
Return(int64(1), nil).
Return([]store.DeletedDevice{{UID: "device-1", TenantID: "tenant-1", Status: models.DeviceStatusRemoved}}, nil).
Once()
queryOptionsMock.
On("Match", mock.MatchedBy(matchFilter())).
Expand All @@ -339,7 +385,7 @@ func TestService_DeviceCleanup(t *testing.T) {
Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-2"}).
Return(int64(1), nil).
Return([]store.DeletedDevice{{UID: "device-2", TenantID: "tenant-2", Status: models.DeviceStatusRemoved}}, nil).
Once()
queryOptionsMock.
On("Match", mock.MatchedBy(matchFilter())).
Expand Down Expand Up @@ -385,6 +431,92 @@ func TestService_DeviceCleanup(t *testing.T) {
}
}

func TestService_EphemeralCleanup(t *testing.T) {
ctx := context.Background()

expired := []models.Device{
{UID: "device-1", TenantID: "tenant-1", Status: models.DeviceStatusAccepted},
{UID: "device-2", TenantID: "tenant-1", Status: models.DeviceStatusPending},
{UID: "device-3", TenantID: "tenant-2", Status: models.DeviceStatusAccepted},
}

cases := []struct {
description string
requiredMocks func(storeMock *storemock.MockStore)
expected error
}{
{
description: "fails when cannot delete the devices",
requiredMocks: func(storeMock *storemock.MockStore) {
storeMock.On("DeviceListExpiredEphemeral", ctx).Return(expired, nil).Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2", "device-3"}).
Return(nil, errors.New("delete error")).
Once()
},
expected: errors.New("delete error"),
},
{
description: "decrements each namespace's counter for the status of every device it deleted",
requiredMocks: func(storeMock *storemock.MockStore) {
storeMock.On("DeviceListExpiredEphemeral", ctx).Return(expired, nil).Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2", "device-3"}).
Return([]store.DeletedDevice{
{UID: "device-1", TenantID: "tenant-1", Status: models.DeviceStatusAccepted},
{UID: "device-2", TenantID: "tenant-1", Status: models.DeviceStatusPending},
{UID: "device-3", TenantID: "tenant-2", Status: models.DeviceStatusAccepted},
}, nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-1"), models.DeviceStatusAccepted, int64(-1)).
Return(nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-1"), models.DeviceStatusPending, int64(-1)).
Return(nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-2"), models.DeviceStatusAccepted, int64(-1)).
Return(nil).
Once()
},
expected: nil,
},
{
description: "decrements the counters only by the devices it deleted",
requiredMocks: func(storeMock *storemock.MockStore) {
storeMock.On("DeviceListExpiredEphemeral", ctx).Return(expired, nil).Once()
storeMock.
On("DeviceDeleteMany", ctx, []string{"device-1", "device-2", "device-3"}).
Return([]store.DeletedDevice{
{UID: "device-3", TenantID: "tenant-2", Status: models.DeviceStatusAccepted},
}, nil).
Once()
storeMock.
On("NamespaceIncrementDeviceCount", ctx, scope.MustBounded("tenant-2"), models.DeviceStatusAccepted, int64(-1)).
Return(nil).
Once()
},
expected: nil,
},
}

for _, tc := range cases {
t.Run(tc.description, func(tt *testing.T) {
storeMock := storemock.NewMockStore(tt)
storeMock.
On("WithTransaction", ctx, mock.AnythingOfType("store.TransactionCb")).
Return(func(ctx context.Context, cb store.TransactionCb) error { return cb(ctx) }).
Once()
tc.requiredMocks(storeMock)

s := NewService(storeMock, privateKey, publicKey, cache.NewNullCache())
require.Equal(tt, tc.expected, s.EphemeralCleanup()(ctx))
})
}
}

func TestService_NamespaceDeviceCountSync(t *testing.T) {
storeMock := storemock.NewMockStore(t)

Expand Down
15 changes: 13 additions & 2 deletions server/api/store/device.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,14 @@ import (
"github.com/shellhub-io/shellhub/pkg/models"
)

// DeletedDevice is a device [DeviceStore.DeviceDeleteMany] deleted, with the namespace and status
// it had, which are what a caller adjusts the namespace's device counts by.
type DeletedDevice struct {
UID string
TenantID string
Status models.DeviceStatus
}

// DeviceAcceptable selects how the acceptable flag is computed for the devices a query
// returns. It depends on the namespace's limit, so it cannot be a stored column.
type DeviceAcceptable uint
Expand Down Expand Up @@ -108,8 +116,11 @@ type DeviceStore interface {
DeviceLockMAC(ctx context.Context, sc scope.Scope, mac string) error

DeviceDelete(ctx context.Context, device *models.Device) error
// DeviceDeleteMany deletes multiple devices by their UIDs.
DeviceDeleteMany(ctx context.Context, uids []string) (deletedCount int64, err error)
// DeviceDeleteMany deletes the devices uids names, with their sessions, and returns the ones it
// deleted. A uid that names no device, including one another call deleted first, is left out,
// so the result can be shorter than uids, and an empty uids deletes nothing. It returns ctx's
// error when ctx ends, and an error wrapping [ErrInternal] for any other database failure.
DeviceDeleteMany(ctx context.Context, uids []string) (deleted []DeletedDevice, err error)
Comment thread
otavio marked this conversation as resolved.
Comment thread
otavio marked this conversation as resolved.

// DeviceListExpiredEphemeral lists ephemeral devices that have stayed offline longer than their
// per-device ephemeral timeout (in minutes), so the cleanup can remove them.
Expand Down
18 changes: 10 additions & 8 deletions server/api/store/mocks/mock_store.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading