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 { type ephemeralPeer struct {
id string id string
account *Account accountID string
deadline time.Time deadline time.Time
next *ephemeralPeer next *ephemeralPeer
} }
// todo: consider to remove peer from ephemeral list when the peer has been deleted via API. If we do not do it // 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) 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() e.peersLock.Lock()
defer e.peersLock.Unlock() defer e.peersLock.Unlock()
@ -117,7 +111,7 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.
return return
} }
e.addPeer(peer.ID, a, newDeadLine()) e.addPeer(peer.AccountID, peer.ID, newDeadLine())
if e.timer == nil { if e.timer == nil {
e.timer = time.AfterFunc(e.headPeer.deadline.Sub(timeNow()), func() { e.timer = time.AfterFunc(e.headPeer.deadline.Sub(timeNow()), func() {
e.cleanup(ctx) e.cleanup(ctx)
@ -126,17 +120,21 @@ func (e *EphemeralManager) OnPeerDisconnected(ctx context.Context, peer *nbpeer.
} }
func (e *EphemeralManager) loadEphemeralPeers(ctx context.Context) { 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() t := newDeadLine()
count := 0 count := 0
for _, a := range accounts { for _, p := range peers {
for id, p := range a.Peers { if p.Ephemeral {
if p.Ephemeral { count++
count++ e.addPeer(p.AccountID, p.ID, t)
e.addPeer(id, a, t)
}
} }
} }
log.WithContext(ctx).Debugf("loaded ephemeral peer(s): %d", count) 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 { for id, p := range deletePeers {
log.WithContext(ctx).Debugf("delete ephemeral peer: %s", id) 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 { if err != nil {
log.WithContext(ctx).Errorf("failed to delete ephemeral peer: %s", err) 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{ ep := &ephemeralPeer{
id: id, id: peerID,
account: account, accountID: accountID,
deadline: deadline, deadline: deadline,
} }
if e.headPeer == nil { if e.headPeer == nil {

View File

@ -7,7 +7,6 @@ import (
"time" "time"
nbpeer "github.com/netbirdio/netbird/management/server/peer" nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/status"
) )
type MockStore struct { type MockStore struct {
@ -15,17 +14,14 @@ type MockStore struct {
account *Account account *Account
} }
func (s *MockStore) GetAllAccounts(_ context.Context) []*Account { func (s *MockStore) GetAllEphemeralPeers(_ context.Context, _ LockingStrength) ([]*nbpeer.Peer, error) {
return []*Account{s.account} var peers []*nbpeer.Peer
} for _, v := range s.account.Peers {
if v.Ephemeral {
func (s *MockStore) GetAccountByPeerID(_ context.Context, peerId string) (*Account, error) { peers = append(peers, v)
_, ok := s.account.Peers[peerId] }
if ok {
return s.account, nil
} }
return peers, nil
return nil, status.NewPeerNotFoundError(peerId)
} }
type MocAccountManager struct { type MocAccountManager struct {