Refactor ephemeral peers

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
This commit is contained in:
bcmmbaga 2024-11-14 13:04:49 +03:00
parent 6315644065
commit 8420a52563
No known key found for this signature in database
GPG Key ID: 511EED5C928AD547
2 changed files with 28 additions and 34 deletions

View File

@ -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 {

View File

@ -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 {