forked from pingcap/tiflow
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
pkg/p2p: add some helper structs for integrating p2p (pingcap#65)
- Loading branch information
Showing
5 changed files
with
489 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,149 @@ | ||
| package p2p | ||
|
|
||
| import ( | ||
| "context" | ||
| "sync" | ||
| "time" | ||
|
|
||
| "github.com/pingcap/errors" | ||
| "github.com/pingcap/tiflow/dm/pkg/log" | ||
| "go.uber.org/atomic" | ||
| "go.uber.org/zap" | ||
| ) | ||
|
|
||
| const ( | ||
| defaultHandlerOperationTimeout = 3 * time.Second | ||
| ) | ||
|
|
||
| // MessageHandlerManager is for managing message topic handlers. | ||
| // NOTE: for each topic, only one handler is allowed. | ||
| type MessageHandlerManager interface { | ||
| RegisterHandler(ctx context.Context, topic Topic, tpi TypeInformation, fn HandlerFunc) (bool, error) | ||
|
|
||
| UnregisterHandler(ctx context.Context, topic Topic) (bool, error) | ||
|
|
||
| CheckError(ctx context.Context) error | ||
|
|
||
| // Clean unregisters all existing handlers. | ||
| Clean(ctx context.Context) error | ||
|
|
||
| // SetTimeout sets the timeout for handler operations. | ||
| // A timeout is needed because the underlying handler operations are | ||
| // asynchronous in the MessageServer. | ||
| SetTimeout(timeout time.Duration) | ||
| } | ||
|
|
||
| func newMessageHandlerManager(registrar handlerRegistrar) MessageHandlerManager { | ||
| return &messageHandlerManagerImpl{ | ||
| messageServer: registrar, | ||
| timeout: atomic.NewDuration(defaultHandlerOperationTimeout), | ||
| topics: make(map[Topic]<-chan error), | ||
| } | ||
| } | ||
|
|
||
| // handlerRegistrar is an interface for the handler management-related | ||
| // functionalities of MessageServer. | ||
| // This interface is for easier unit-testing. | ||
| type handlerRegistrar interface { | ||
| SyncAddHandler(context.Context, Topic, TypeInformation, HandlerFunc) (<-chan error, error) | ||
| SyncRemoveHandler(context.Context, Topic) error | ||
| } | ||
|
|
||
| type messageHandlerManagerImpl struct { | ||
| messageServer handlerRegistrar | ||
| // timeout is atomic to avoid unnecessary blocking | ||
| timeout *atomic.Duration | ||
|
|
||
| // mu protects topics | ||
| mu sync.Mutex | ||
| topics map[Topic]<-chan error | ||
| } | ||
|
|
||
| func (m *messageHandlerManagerImpl) RegisterHandler( | ||
| ctx context.Context, | ||
| topic Topic, | ||
| tpi TypeInformation, | ||
| fn HandlerFunc, | ||
| ) (bool, error) { | ||
| m.mu.Lock() | ||
| defer m.mu.Unlock() | ||
|
|
||
| if _, ok := m.topics[topic]; ok { | ||
| // A handler for this topic already exists. | ||
| return false, nil | ||
| } | ||
|
|
||
| ctx, cancel := m.makeContext(ctx) | ||
| defer cancel() | ||
|
|
||
| errCh, err := m.messageServer.SyncAddHandler(ctx, topic, tpi, fn) | ||
| if err != nil { | ||
| return false, errors.Trace(err) | ||
| } | ||
| m.topics[topic] = errCh | ||
|
|
||
| return true, nil | ||
| } | ||
|
|
||
| func (m *messageHandlerManagerImpl) UnregisterHandler(ctx context.Context, topic Topic) (bool, error) { | ||
| m.mu.Lock() | ||
| defer m.mu.Unlock() | ||
|
|
||
| if _, ok := m.topics[topic]; !ok { | ||
| // The handler for this topic does not exist | ||
| return false, nil | ||
| } | ||
|
|
||
| ctx, cancel := m.makeContext(ctx) | ||
| defer cancel() | ||
|
|
||
| if err := m.messageServer.SyncRemoveHandler(ctx, topic); err != nil { | ||
| return false, errors.Trace(err) | ||
| } | ||
| delete(m.topics, topic) | ||
|
|
||
| return true, nil | ||
| } | ||
|
|
||
| func (m *messageHandlerManagerImpl) CheckError(ctx context.Context) error { | ||
| m.mu.Lock() | ||
| defer m.mu.Unlock() | ||
|
|
||
| for topic, errCh := range m.topics { | ||
| select { | ||
| case <-ctx.Done(): | ||
| return errors.Trace(ctx.Err()) | ||
| case err := <-errCh: | ||
| if err == nil { | ||
| continue | ||
| } | ||
| log.L().Warn("handler error received", | ||
| zap.String("topic", topic)) | ||
| return errors.Trace(err) | ||
| default: | ||
| } | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| func (m *messageHandlerManagerImpl) Clean(ctx context.Context) error { | ||
| m.mu.Lock() | ||
| defer m.mu.Unlock() | ||
|
|
||
| for topic := range m.topics { | ||
| if err := m.messageServer.SyncRemoveHandler(ctx, topic); err != nil { | ||
| return errors.Trace(err) | ||
| } | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
|
|
||
| func (m *messageHandlerManagerImpl) SetTimeout(timeout time.Duration) { | ||
| m.timeout.Store(timeout) | ||
| } | ||
|
|
||
| func (m *messageHandlerManagerImpl) makeContext(parent context.Context) (context.Context, context.CancelFunc) { | ||
| timeout := m.timeout.Load() | ||
| return context.WithTimeout(parent, timeout) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,125 @@ | ||
| package p2p | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
| "time" | ||
|
|
||
| "github.com/pingcap/errors" | ||
| p2pImpl "github.com/pingcap/tiflow/pkg/p2p" | ||
| "github.com/stretchr/testify/mock" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| // handlerRegistrar must be implemented by MessageServer | ||
| var _ handlerRegistrar = (*p2pImpl.MessageServer)(nil) | ||
|
|
||
| type mockHandlerRegistrar struct { | ||
| mock.Mock | ||
| } | ||
|
|
||
| func (r *mockHandlerRegistrar) SyncAddHandler( | ||
| ctx context.Context, | ||
| topic Topic, | ||
| information TypeInformation, | ||
| handlerFunc HandlerFunc, | ||
| ) (<-chan error, error) { | ||
| args := r.Called(ctx, topic, information, handlerFunc) | ||
| return args.Get(0).(<-chan error), args.Error(1) | ||
| } | ||
|
|
||
| func (r *mockHandlerRegistrar) SyncRemoveHandler(ctx context.Context, topic Topic) error { | ||
| args := r.Called(ctx, topic) | ||
| return args.Error(0) | ||
| } | ||
|
|
||
| type msgContent struct{} | ||
|
|
||
| func TestMessageHandlerManagerBasics(t *testing.T) { | ||
| ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) | ||
| defer cancel() | ||
|
|
||
| registrar := &mockHandlerRegistrar{} | ||
| manager := newMessageHandlerManager(registrar) | ||
|
|
||
| errCh1 := make(chan error, 1) | ||
| registrar.On("SyncAddHandler", mock.Anything, "test-topic-1", &msgContent{}, mock.Anything). | ||
| Return((<-chan error)(errCh1), nil) | ||
| ok, err := manager.RegisterHandler(ctx, "test-topic-1", &msgContent{}, func(NodeID, MessageValue) error { | ||
| // This function does not matter here | ||
| return nil | ||
| }) | ||
| require.NoError(t, err) | ||
| require.True(t, ok) | ||
| registrar.AssertExpectations(t) | ||
|
|
||
| // Test duplicate handler | ||
| ok, err = manager.RegisterHandler(ctx, "test-topic-1", &msgContent{}, func(NodeID, MessageValue) error { | ||
| // This function does not matter here | ||
| return nil | ||
| }) | ||
| require.NoError(t, err) | ||
| require.False(t, ok) | ||
|
|
||
| errCh2 := make(chan error, 1) | ||
| registrar.ExpectedCalls = nil | ||
| registrar.On("SyncAddHandler", mock.Anything, "test-topic-2", &msgContent{}, mock.Anything). | ||
| Return((<-chan error)(errCh2), nil) | ||
| ok, err = manager.RegisterHandler(ctx, "test-topic-2", &msgContent{}, func(NodeID, MessageValue) error { | ||
| // This function does not matter here | ||
| return nil | ||
| }) | ||
| require.NoError(t, err) | ||
| require.True(t, ok) | ||
| registrar.AssertExpectations(t) | ||
|
|
||
| err = manager.CheckError(ctx) | ||
| require.NoError(t, err) | ||
|
|
||
| errCh1 <- errors.New("fake error") | ||
| err = manager.CheckError(ctx) | ||
| require.Error(t, err) | ||
|
|
||
| registrar.ExpectedCalls = nil | ||
| registrar.On("SyncRemoveHandler", mock.Anything, "test-topic-1").Return(nil) | ||
| ok, err = manager.UnregisterHandler(ctx, "test-topic-1") | ||
| require.NoError(t, err) | ||
| require.True(t, ok) | ||
| registrar.AssertExpectations(t) | ||
|
|
||
| // duplicate unregister | ||
| ok, err = manager.UnregisterHandler(ctx, "test-topic-1") | ||
| require.NoError(t, err) | ||
| require.False(t, ok) | ||
|
|
||
| registrar.ExpectedCalls = nil | ||
| registrar.On("SyncRemoveHandler", mock.Anything, "test-topic-2").Return(nil) | ||
| err = manager.Clean(ctx) | ||
| require.NoError(t, err) | ||
| } | ||
|
|
||
| func TestMessageHandlerManagerTimeout(t *testing.T) { | ||
| ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) | ||
| defer cancel() | ||
|
|
||
| registrar := &mockHandlerRegistrar{} | ||
| manager := newMessageHandlerManager(registrar) | ||
| manager.SetTimeout(time.Duration(0)) | ||
|
|
||
| registrar.On("SyncAddHandler", mock.Anything, "test-topic-1", &msgContent{}, mock.Anything). | ||
| Return((<-chan error)(nil), errors.New("fake error")). | ||
| Run(func(args mock.Arguments) { | ||
| ctx := args.Get(0).(context.Context) | ||
| select { | ||
| case <-ctx.Done(): | ||
| default: | ||
| require.Fail(t, "context should have been canceled") | ||
| } | ||
| }) | ||
|
|
||
| _, err := manager.RegisterHandler(ctx, "test-topic-1", &msgContent{}, func(NodeID, MessageValue) error { | ||
| // This function does not matter here | ||
| return nil | ||
| }) | ||
| require.Error(t, err) | ||
| } |
Oops, something went wrong.