diff --git a/management/server/ephemeral.go b/management/server/ephemeral.go index 590b1d708..6e245ec5a 100644 --- a/management/server/ephemeral.go +++ b/management/server/ephemeral.go @@ -20,10 +20,10 @@ var ( ) type ephemeralPeer struct { - id string - account *Account - deadline time.Time - next *ephemeralPeer + id string + accountID string + deadline time.Time + next *ephemeralPeer } // todo: consider to remove peer from ephemeral list when the peer has been deleted via API. If we do not do it @@ -104,12 +104,6 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer. log.WithContext(ctx).Tracef("add peer to ephemeral list: %s", peer.ID) - a, err := e.store.GetAccountByPeerID(context.Background(), peer.ID) - if err != nil { - log.WithContext(ctx).Errorf("failed to add peer to ephemeral list: %s", err) - return - } - e.peersLock.Lock() defer e.peersLock.Unlock() @@ -117,7 +111,7 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer. return } - e.addPeer(peer.ID, a, newDeadLine()) + e.addPeer(peer.AccountID, peer.ID, newDeadLine()) if e.timer == nil { e.timer = time.AfterFunc(e.headPeer.deadline.Sub(timeNow()), func() { e.cleanup(ctx) @@ -126,17 +120,21 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer. } func (e *EphemeralManager) loadEphemeralPeers(ctx context.Context) { - accounts := e.store.GetAllAccounts(context.Background()) + peers, err := e.store.GetAllEphemeralPeers(ctx, LockingStrengthShare) + if err != nil { + log.WithContext(ctx).Debugf("failed to load ephemeral peers: %s", err) + return + } + t := newDeadLine() count := 0 - for _, a := range accounts { - for id, p := range a.Peers { - if p.Ephemeral { - count++ - e.addPeer(id, a, t) - } + for _, p := range peers { + if p.Ephemeral { + count++ + e.addPeer(p.AccountID, p.ID, t) } } + log.WithContext(ctx).Debugf("loaded ephemeral peer(s): %d", count) } @@ -170,18 +168,18 @@ func (e *EphemeralManager) cleanup(ctx context.Context) { for id, p := range deletePeers { log.WithContext(ctx).Debugf("delete ephemeral peer: %s", id) - err := e.accountManager.DeletePeer(ctx, p.account.Id, id, activity.SystemInitiator) + err := e.accountManager.DeletePeer(ctx, p.accountID, id, activity.SystemInitiator) if err != nil { log.WithContext(ctx).Errorf("failed to delete ephemeral peer: %s", err) } } } -func (e *EphemeralManager) addPeer(id string, account *Account, deadline time.Time) { +func (e *EphemeralManager) addPeer(accountID string, peerID string, deadline time.Time) { ep := &ephemeralPeer{ - id: id, - account: account, - deadline: deadline, + id: peerID, + accountID: accountID, + deadline: deadline, } if e.headPeer == nil { diff --git a/management/server/ephemeral_test.go b/management/server/ephemeral_test.go index 1390352a5..00e5d777a 100644 --- a/management/server/ephemeral_test.go +++ b/management/server/ephemeral_test.go @@ -7,7 +7,6 @@ import ( "time" nbpeer "github.com/netbirdio/netbird/management/server/peer" - "github.com/netbirdio/netbird/management/server/status" ) type MockStore struct { @@ -15,17 +14,14 @@ type MockStore struct { account *Account } -func (s *MockStore) GetAllAccounts(_ context.Context) []*Account { - return []*Account{s.account} -} - -func (s *MockStore) GetAccountByPeerID(_ context.Context, peerId string) (*Account, error) { - _, ok := s.account.Peers[peerId] - if ok { - return s.account, nil +func (s *MockStore) GetAllEphemeralPeers(_ context.Context, _ LockingStrength) ([]*nbpeer.Peer, error) { + var peers []*nbpeer.Peer + for _, v := range s.account.Peers { + if v.Ephemeral { + peers = append(peers, v) + } } - - return nil, status.NewPeerNotFoundError(peerId) + return peers, nil } type MocAccountManager struct {