From f422c4165437679a1470893e7917a8ef597a2277 Mon Sep 17 00:00:00 2001 From: dmitri-netbird Date: Fri, 11 Sep 2026 17:07:26 +0200 Subject: [PATCH] [management] extract peer update logic and wrap it in tests (#7338) * extract peer update loop into a dedicated struct and wrap it in tests Signed-off-by: Dmitri Dolguikh * make linter happy Signed-off-by: Dmitri Dolguikh --------- Signed-off-by: Dmitri Dolguikh --- encryption/message.go | 10 ++ .../shared/grpc/peer_update_handler.go | 135 +++++++++++++++ .../shared/grpc/peer_update_handler_test.go | 155 ++++++++++++++++++ management/internals/shared/grpc/server.go | 88 +--------- .../internals/shared/grpc/sync_sender_mock.go | 70 ++++++++ management/internals/shared/grpc/token_mgr.go | 2 + .../internals/shared/grpc/token_mgr_mock.go | 111 +++++++++++++ .../internals/shared/grpc/update_debouncer.go | 8 + .../shared/grpc/update_debouncer_mock.go | 96 +++++++++++ 9 files changed, 589 insertions(+), 86 deletions(-) create mode 100644 management/internals/shared/grpc/peer_update_handler.go create mode 100644 management/internals/shared/grpc/peer_update_handler_test.go create mode 100644 management/internals/shared/grpc/sync_sender_mock.go create mode 100644 management/internals/shared/grpc/token_mgr_mock.go create mode 100644 management/internals/shared/grpc/update_debouncer_mock.go diff --git a/encryption/message.go b/encryption/message.go index 6e4cd7391..2bf2c59dc 100644 --- a/encryption/message.go +++ b/encryption/message.go @@ -6,6 +6,16 @@ import ( "golang.zx2c4.com/wireguard/wgctrl/wgtypes" ) +type Encrypter interface { + EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) +} + +type DefaultEncrypter struct{} + +func (e DefaultEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) { + return EncryptMessage(remotePubKey, ourPrivateKey, message) +} + // EncryptMessage encrypts a body of the given protobuf Message func EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) { byteResp, err := pb.Marshal(message) diff --git a/management/internals/shared/grpc/peer_update_handler.go b/management/internals/shared/grpc/peer_update_handler.go new file mode 100644 index 000000000..d2f403841 --- /dev/null +++ b/management/internals/shared/grpc/peer_update_handler.go @@ -0,0 +1,135 @@ +package grpc + +import ( + "context" + "time" + + "github.com/netbirdio/netbird/encryption" + "github.com/netbirdio/netbird/management/internals/controllers/network_map" + "github.com/netbirdio/netbird/management/server/telemetry" + "github.com/netbirdio/netbird/shared/management/proto" + log "github.com/sirupsen/logrus" + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func PeerUpdateHandlerFactory( + peerKey wgtypes.Key, + updates chan *network_map.UpdateMessage, + secretsManager SecretsManager, + srv proto.ManagementService_SyncServer, + cleanupfunc func()) *PeerUpdateHandler { + return &PeerUpdateHandler{ + peerKey: peerKey, + updates: updates, + secretsManager: secretsManager, + srv: srv, + encrypter: encryption.DefaultEncrypter{}, + debouncer: NewUpdateDebouncer(1000 * time.Millisecond), + cleanupFunc: cleanupfunc, + } +} + +// PeerUpdateHandler sends updates to the connected peer until the updates channel is closed. +// It implements a backpressure mechanism that sends the first update immediately, +// then debounces subsequent rapid updates, ensuring only the latest update is sent +// after a quiet period. +type PeerUpdateHandler struct { + peerKey wgtypes.Key + updates chan *network_map.UpdateMessage + appMetrics telemetry.AppMetrics + secretsManager SecretsManager + srv syncSender + encrypter encryption.Encrypter + debouncer Debouncer + cleanupFunc func() +} + +func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler { + pu.appMetrics = appMetrics + return pu +} + +//go:generate go tool mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc +type syncSender interface { + Send(*proto.EncryptedMessage) error + Context() context.Context +} + +func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error { + log.WithContext(ctx).Tracef("starting to handle updates for peer %s", pu.peerKey.String()) + + defer pu.debouncer.Stop() + + for { + select { + // condition when there are some updates + // todo set the updates channel size to 1 + case update, open := <-pu.updates: + if pu.appMetrics != nil { + pu.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(pu.updates) + 1) + } + + if !open { + log.WithContext(ctx).Debugf("updates channel for peer %s was closed", pu.peerKey.String()) + pu.cleanupFunc() + return nil + } + + log.WithContext(ctx).Tracef("received an update for peer %s", pu.peerKey.String()) + if pu.debouncer.ProcessUpdate(update) { + // Send immediately (first update or after quiet period) + if err := pu.SendUpdate(ctx, update); err != nil { + log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err) + return err + } + } + + // Timer expired - quiet period reached, send pending updates if any + case <-pu.debouncer.TimerChannel(): + pendingUpdates := pu.debouncer.GetPendingUpdates() + if len(pendingUpdates) == 0 { + continue + } + log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), pu.peerKey.String()) + for _, pendingUpdate := range pendingUpdates { + if err := pu.SendUpdate(ctx, pendingUpdate); err != nil { + log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err) + return err + } + } + + // condition when client <-> server connection has been terminated + case <-pu.srv.Context().Done(): + // happens when connection drops, e.g. client disconnects + log.WithContext(ctx).Debugf("stream of peer %s has been closed", pu.peerKey.String()) + pu.cleanupFunc() + return pu.srv.Context().Err() + } + } +} + +func (pu *PeerUpdateHandler) SendUpdate(ctx context.Context, update *network_map.UpdateMessage) error { + key, err := pu.secretsManager.GetWGKey() + if err != nil { + pu.cleanupFunc() + return status.Errorf(codes.Internal, "failed processing update message") + } + + encryptedResp, err := pu.encrypter.EncryptMessage(pu.peerKey, key, update.Update) + if err != nil { + pu.cleanupFunc() + return status.Errorf(codes.Internal, "failed processing update message") + } + err = pu.srv.Send(&proto.EncryptedMessage{ + WgPubKey: key.PublicKey().String(), + Body: encryptedResp, + }) + if err != nil { + pu.cleanupFunc() + return status.Errorf(codes.Internal, "failed sending update message") + } + log.WithContext(ctx).Tracef("sent an update to peer %s", pu.peerKey.String()) + return nil +} diff --git a/management/internals/shared/grpc/peer_update_handler_test.go b/management/internals/shared/grpc/peer_update_handler_test.go new file mode 100644 index 000000000..02de49c47 --- /dev/null +++ b/management/internals/shared/grpc/peer_update_handler_test.go @@ -0,0 +1,155 @@ +package grpc + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + pb "github.com/golang/protobuf/proto" //nolint + "github.com/netbirdio/netbird/management/internals/controllers/network_map" + "github.com/netbirdio/netbird/shared/management/proto" + "github.com/stretchr/testify/assert" + "go.uber.org/mock/gomock" + "golang.zx2c4.com/wireguard/wgctrl/wgtypes" +) + +func TestSendPeerUpdates_FirstUpdate(t *testing.T) { + ctrl := gomock.NewController(t) + secretsManager := NewMockSecretsManager(ctrl) + updateDebouncer := NewMockDebouncer(ctrl) + syncSender := NewMocksyncSender(ctrl) + + pu := PeerUpdateHandler{ + peerKey: mustGenerateKey(t), + updates: make(chan *network_map.UpdateMessage), + secretsManager: secretsManager, + encrypter: testEncrypter{}, + debouncer: updateDebouncer, + srv: syncSender, + cleanupFunc: func() {}, + } + + msg := network_map.UpdateMessage{ + Update: &proto.SyncResponse{Version: 1}, + } + + timeCh := make(chan time.Time) + srvCtx := context.TODO() + srvKey := mustGenerateKey(t) + // mock a first update, should send it right away + updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true) + updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh) + syncSender.EXPECT().Context().AnyTimes().Return(srvCtx) + secretsManager.EXPECT().GetWGKey().Return(srvKey, nil) + syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}}) + updateDebouncer.EXPECT().Stop() + + var wg sync.WaitGroup + wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck + pu.updates <- &msg + close(pu.updates) + wg.Wait() +} + +func TestSendPeerUpdates_TimerUpdate(t *testing.T) { + ctrl := gomock.NewController(t) + secretsManager := NewMockSecretsManager(ctrl) + updateDebouncer := NewMockDebouncer(ctrl) + syncSender := NewMocksyncSender(ctrl) + + pu := PeerUpdateHandler{ + peerKey: mustGenerateKey(t), + updates: make(chan *network_map.UpdateMessage), + secretsManager: secretsManager, + encrypter: testEncrypter{}, + debouncer: updateDebouncer, + srv: syncSender, + cleanupFunc: func() {}, + } + + msg := network_map.UpdateMessage{ + Update: &proto.SyncResponse{Version: 1}, + } + + timeCh := make(chan time.Time) + srvCtx := context.TODO() + srvKey := mustGenerateKey(t) + updateDebouncer.EXPECT().GetPendingUpdates().Return([]*network_map.UpdateMessage{&msg}) + updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh) + syncSender.EXPECT().Context().AnyTimes().Return(srvCtx) + secretsManager.EXPECT().GetWGKey().Return(srvKey, nil) + syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}}) + updateDebouncer.EXPECT().Stop() + + var wg sync.WaitGroup + wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck + timeCh <- time.Now() + close(pu.updates) + wg.Wait() +} + +func TestSendPeerUpdates_ServerContextDone(t *testing.T) { + ctrl := gomock.NewController(t) + secretsManager := NewMockSecretsManager(ctrl) + updateDebouncer := NewMockDebouncer(ctrl) + syncSender := NewMocksyncSender(ctrl) + + pu := PeerUpdateHandler{ + peerKey: mustGenerateKey(t), + updates: make(chan *network_map.UpdateMessage), + secretsManager: secretsManager, + encrypter: testEncrypter{}, + debouncer: updateDebouncer, + srv: syncSender, + cleanupFunc: func() {}, + } + + timeCh := make(chan time.Time) + srvCtx, cancel := context.WithCancel(context.TODO()) + updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh) + syncSender.EXPECT().Context().AnyTimes().Return(srvCtx) + updateDebouncer.EXPECT().Stop() + + var wg sync.WaitGroup + wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck + cancel() + wg.Wait() +} + +func mustGenerateKey(t *testing.T) wgtypes.Key { + t.Helper() + k, err := wgtypes.GenerateKey() + assert.NoError(t, err) + return k +} + +func mustMarshal(t *testing.T, msg *network_map.UpdateMessage) []byte { + t.Helper() + r, err := pb.Marshal(msg.Update) + assert.NoError(t, err) + return r +} + +type testEncrypter struct{} + +func (testEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) { + return pb.Marshal(message) +} + +type pbMatcher struct { + x pb.Message +} + +func (pbm pbMatcher) Matches(x any) bool { + msg, ok := x.(pb.Message) + if !ok { + return false + } + return pb.Equal(pbm.x, msg) +} + +func (pbm pbMatcher) String() string { + return fmt.Sprintf("is equal to %s (%T)", pbm.x, pbm.x) +} diff --git a/management/internals/shared/grpc/server.go b/management/internals/shared/grpc/server.go index a9cc0ad36..c178b6fa1 100644 --- a/management/internals/shared/grpc/server.go +++ b/management/internals/shared/grpc/server.go @@ -337,7 +337,8 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S s.syncSem.Add(-1) - return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv, syncStart) + return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart) }). + WithMetrics(s.appMetrics).HandleUpdates(ctx) } func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementService_JobServer) (wgtypes.Key, error) { @@ -404,91 +405,6 @@ func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgt } } -// handleUpdates sends updates to the connected peer until the updates channel is closed. -// It implements a backpressure mechanism that sends the first update immediately, -// then debounces subsequent rapid updates, ensuring only the latest update is sent -// after a quiet period. -func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error { - log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String()) - - // Create a debouncer for this peer connection - debouncer := NewUpdateDebouncer(1000 * time.Millisecond) - defer debouncer.Stop() - - for { - select { - // condition when there are some updates - // todo set the updates channel size to 1 - case update, open := <-updates: - if s.appMetrics != nil { - s.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(updates) + 1) - } - - if !open { - log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String()) - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) - return nil - } - - log.WithContext(ctx).Tracef("received an update for peer %s", peerKey.String()) - if debouncer.ProcessUpdate(update) { - // Send immediately (first update or after quiet period) - if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime); err != nil { - log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err) - return err - } - } - - // Timer expired - quiet period reached, send pending updates if any - case <-debouncer.TimerChannel(): - pendingUpdates := debouncer.GetPendingUpdates() - if len(pendingUpdates) == 0 { - continue - } - log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), peerKey.String()) - for _, pendingUpdate := range pendingUpdates { - if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime); err != nil { - log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err) - return err - } - } - - // condition when client <-> server connection has been terminated - case <-srv.Context().Done(): - // happens when connection drops, e.g. client disconnects - log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String()) - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) - return srv.Context().Err() - } - } -} - -// sendUpdate encrypts the update message using the peer key and the server's wireguard key, -// then sends the encrypted message to the connected peer via the sync server. -func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error { - key, err := s.secretsManager.GetWGKey() - if err != nil { - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) - return status.Errorf(codes.Internal, "failed processing update message") - } - - encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update) - if err != nil { - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) - return status.Errorf(codes.Internal, "failed processing update message") - } - err = srv.Send(&proto.EncryptedMessage{ - WgPubKey: key.PublicKey().String(), - Body: encryptedResp, - }) - if err != nil { - s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime) - return status.Errorf(codes.Internal, "failed sending update message") - } - log.WithContext(ctx).Tracef("sent an update to peer %s", peerKey.String()) - return nil -} - // sendJob encrypts the update message using the peer key and the server's wireguard key, // then sends the encrypted message to the connected peer via the sync server. func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Event, srv proto.ManagementService_JobServer) error { diff --git a/management/internals/shared/grpc/sync_sender_mock.go b/management/internals/shared/grpc/sync_sender_mock.go new file mode 100644 index 000000000..3d1696f59 --- /dev/null +++ b/management/internals/shared/grpc/sync_sender_mock.go @@ -0,0 +1,70 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: ./peer_update_handler.go +// +// Generated by this command: +// +// mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc +// + +// Package grpc is a generated GoMock package. +package grpc + +import ( + context "context" + reflect "reflect" + + proto "github.com/netbirdio/netbird/shared/management/proto" + gomock "go.uber.org/mock/gomock" +) + +// MocksyncSender is a mock of syncSender interface. +type MocksyncSender struct { + ctrl *gomock.Controller + recorder *MocksyncSenderMockRecorder + isgomock struct{} +} + +// MocksyncSenderMockRecorder is the mock recorder for MocksyncSender. +type MocksyncSenderMockRecorder struct { + mock *MocksyncSender +} + +// NewMocksyncSender creates a new mock instance. +func NewMocksyncSender(ctrl *gomock.Controller) *MocksyncSender { + mock := &MocksyncSender{ctrl: ctrl} + mock.recorder = &MocksyncSenderMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MocksyncSender) EXPECT() *MocksyncSenderMockRecorder { + return m.recorder +} + +// Context mocks base method. +func (m *MocksyncSender) Context() context.Context { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Context") + ret0, _ := ret[0].(context.Context) + return ret0 +} + +// Context indicates an expected call of Context. +func (mr *MocksyncSenderMockRecorder) Context() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MocksyncSender)(nil).Context)) +} + +// Send mocks base method. +func (m *MocksyncSender) Send(arg0 *proto.EncryptedMessage) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Send", arg0) + ret0, _ := ret[0].(error) + return ret0 +} + +// Send indicates an expected call of Send. +func (mr *MocksyncSenderMockRecorder) Send(arg0 any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MocksyncSender)(nil).Send), arg0) +} diff --git a/management/internals/shared/grpc/token_mgr.go b/management/internals/shared/grpc/token_mgr.go index fb2d83a9a..4dab5007f 100644 --- a/management/internals/shared/grpc/token_mgr.go +++ b/management/internals/shared/grpc/token_mgr.go @@ -25,6 +25,8 @@ import ( const defaultDuration = 12 * time.Hour // SecretsManager used to manage TURN and relay secrets +// +//go:generate go tool mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc type SecretsManager interface { GenerateTurnToken() (*Token, error) GenerateRelayToken() (*Token, error) diff --git a/management/internals/shared/grpc/token_mgr_mock.go b/management/internals/shared/grpc/token_mgr_mock.go new file mode 100644 index 000000000..e7dde4e3a --- /dev/null +++ b/management/internals/shared/grpc/token_mgr_mock.go @@ -0,0 +1,111 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: ./token_mgr.go +// +// Generated by this command: +// +// mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc +// + +// Package grpc is a generated GoMock package. +package grpc + +import ( + context "context" + reflect "reflect" + + gomock "go.uber.org/mock/gomock" + wgtypes "golang.zx2c4.com/wireguard/wgctrl/wgtypes" +) + +// MockSecretsManager is a mock of SecretsManager interface. +type MockSecretsManager struct { + ctrl *gomock.Controller + recorder *MockSecretsManagerMockRecorder + isgomock struct{} +} + +// MockSecretsManagerMockRecorder is the mock recorder for MockSecretsManager. +type MockSecretsManagerMockRecorder struct { + mock *MockSecretsManager +} + +// NewMockSecretsManager creates a new mock instance. +func NewMockSecretsManager(ctrl *gomock.Controller) *MockSecretsManager { + mock := &MockSecretsManager{ctrl: ctrl} + mock.recorder = &MockSecretsManagerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockSecretsManager) EXPECT() *MockSecretsManagerMockRecorder { + return m.recorder +} + +// CancelRefresh mocks base method. +func (m *MockSecretsManager) CancelRefresh(peerKey string) { + m.ctrl.T.Helper() + m.ctrl.Call(m, "CancelRefresh", peerKey) +} + +// CancelRefresh indicates an expected call of CancelRefresh. +func (mr *MockSecretsManagerMockRecorder) CancelRefresh(peerKey any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRefresh", reflect.TypeOf((*MockSecretsManager)(nil).CancelRefresh), peerKey) +} + +// GenerateRelayToken mocks base method. +func (m *MockSecretsManager) GenerateRelayToken() (*Token, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GenerateRelayToken") + ret0, _ := ret[0].(*Token) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GenerateRelayToken indicates an expected call of GenerateRelayToken. +func (mr *MockSecretsManagerMockRecorder) GenerateRelayToken() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateRelayToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateRelayToken)) +} + +// GenerateTurnToken mocks base method. +func (m *MockSecretsManager) GenerateTurnToken() (*Token, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GenerateTurnToken") + ret0, _ := ret[0].(*Token) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GenerateTurnToken indicates an expected call of GenerateTurnToken. +func (mr *MockSecretsManagerMockRecorder) GenerateTurnToken() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateTurnToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateTurnToken)) +} + +// GetWGKey mocks base method. +func (m *MockSecretsManager) GetWGKey() (wgtypes.Key, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetWGKey") + ret0, _ := ret[0].(wgtypes.Key) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetWGKey indicates an expected call of GetWGKey. +func (mr *MockSecretsManagerMockRecorder) GetWGKey() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWGKey", reflect.TypeOf((*MockSecretsManager)(nil).GetWGKey)) +} + +// SetupRefresh mocks base method. +func (m *MockSecretsManager) SetupRefresh(ctx context.Context, accountID, peerKey string) { + m.ctrl.T.Helper() + m.ctrl.Call(m, "SetupRefresh", ctx, accountID, peerKey) +} + +// SetupRefresh indicates an expected call of SetupRefresh. +func (mr *MockSecretsManagerMockRecorder) SetupRefresh(ctx, accountID, peerKey any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetupRefresh", reflect.TypeOf((*MockSecretsManager)(nil).SetupRefresh), ctx, accountID, peerKey) +} diff --git a/management/internals/shared/grpc/update_debouncer.go b/management/internals/shared/grpc/update_debouncer.go index 8af9c2656..9483007c6 100644 --- a/management/internals/shared/grpc/update_debouncer.go +++ b/management/internals/shared/grpc/update_debouncer.go @@ -6,6 +6,14 @@ import ( "github.com/netbirdio/netbird/management/internals/controllers/network_map" ) +//go:generate go tool mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc +type Debouncer interface { + Stop() + TimerChannel() <-chan time.Time + ProcessUpdate(update *network_map.UpdateMessage) bool + GetPendingUpdates() []*network_map.UpdateMessage +} + // UpdateDebouncer implements a backpressure mechanism that: // - Sends the first update immediately // - Coalesces rapid subsequent network map updates (only latest matters) diff --git a/management/internals/shared/grpc/update_debouncer_mock.go b/management/internals/shared/grpc/update_debouncer_mock.go new file mode 100644 index 000000000..4dff632b7 --- /dev/null +++ b/management/internals/shared/grpc/update_debouncer_mock.go @@ -0,0 +1,96 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: ./update_debouncer.go +// +// Generated by this command: +// +// mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc +// + +// Package grpc is a generated GoMock package. +package grpc + +import ( + reflect "reflect" + time "time" + + network_map "github.com/netbirdio/netbird/management/internals/controllers/network_map" + gomock "go.uber.org/mock/gomock" +) + +// MockDebouncer is a mock of Debouncer interface. +type MockDebouncer struct { + ctrl *gomock.Controller + recorder *MockDebouncerMockRecorder + isgomock struct{} +} + +// MockDebouncerMockRecorder is the mock recorder for MockDebouncer. +type MockDebouncerMockRecorder struct { + mock *MockDebouncer +} + +// NewMockDebouncer creates a new mock instance. +func NewMockDebouncer(ctrl *gomock.Controller) *MockDebouncer { + mock := &MockDebouncer{ctrl: ctrl} + mock.recorder = &MockDebouncerMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockDebouncer) EXPECT() *MockDebouncerMockRecorder { + return m.recorder +} + +// GetPendingUpdates mocks base method. +func (m *MockDebouncer) GetPendingUpdates() []*network_map.UpdateMessage { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetPendingUpdates") + ret0, _ := ret[0].([]*network_map.UpdateMessage) + return ret0 +} + +// GetPendingUpdates indicates an expected call of GetPendingUpdates. +func (mr *MockDebouncerMockRecorder) GetPendingUpdates() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPendingUpdates", reflect.TypeOf((*MockDebouncer)(nil).GetPendingUpdates)) +} + +// ProcessUpdate mocks base method. +func (m *MockDebouncer) ProcessUpdate(update *network_map.UpdateMessage) bool { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ProcessUpdate", update) + ret0, _ := ret[0].(bool) + return ret0 +} + +// ProcessUpdate indicates an expected call of ProcessUpdate. +func (mr *MockDebouncerMockRecorder) ProcessUpdate(update any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProcessUpdate", reflect.TypeOf((*MockDebouncer)(nil).ProcessUpdate), update) +} + +// Stop mocks base method. +func (m *MockDebouncer) Stop() { + m.ctrl.T.Helper() + m.ctrl.Call(m, "Stop") +} + +// Stop indicates an expected call of Stop. +func (mr *MockDebouncerMockRecorder) Stop() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockDebouncer)(nil).Stop)) +} + +// TimerChannel mocks base method. +func (m *MockDebouncer) TimerChannel() <-chan time.Time { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "TimerChannel") + ret0, _ := ret[0].(<-chan time.Time) + return ret0 +} + +// TimerChannel indicates an expected call of TimerChannel. +func (mr *MockDebouncerMockRecorder) TimerChannel() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimerChannel", reflect.TypeOf((*MockDebouncer)(nil).TimerChannel)) +}