Compare commits

...

1 Commits

Author SHA1 Message Date
Vladimir Mihailenco 652668236d Add ConnectHook 2021-06-11 17:15:06 +03:00
1 changed files with 30 additions and 2 deletions

View File

@ -29,8 +29,22 @@ type Hook interface {
AfterProcessPipeline(ctx context.Context, cmds []Cmder) error AfterProcessPipeline(ctx context.Context, cmds []Cmder) error
} }
type ConnectHook interface {
BeforeConnect(ctx context.Context, event ConnectEvent) context.Context
AfterConnect(ctx context.Context, event ConnectEvent)
}
type fullHook interface {
Hook
ConnectHook
}
type ConnectEvent struct {
Err error
}
type hooks struct { type hooks struct {
hooks []Hook hooks []fullHook
} }
func (hs *hooks) lock() { func (hs *hooks) lock() {
@ -44,7 +58,11 @@ func (hs hooks) clone() hooks {
} }
func (hs *hooks) AddHook(hook Hook) { func (hs *hooks) AddHook(hook Hook) {
if hook, ok := hook.(fullHook); ok {
hs.hooks = append(hs.hooks, hook) hs.hooks = append(hs.hooks, hook)
} else {
hs.hooks = append(hs.hooks, dummyConnectHook{Hook: hook})
}
} }
func (hs hooks) process( func (hs hooks) process(
@ -132,6 +150,16 @@ func (hs hooks) withContext(ctx context.Context, fn func() error) error {
return fn() return fn()
} }
type dummyConnectHook struct {
Hook
}
func (dummyConnectHook) BeforeConnect(ctx context.Context, event ConnectEvent) context.Context {
return ctx
}
func (dummyConnectHook) AfterConnect(ctx context.Context, event ConnectEvent) {}
//------------------------------------------------------------------------------ //------------------------------------------------------------------------------
type baseClient struct { type baseClient struct {