Add ConnectHook

This commit is contained in:
Vladimir Mihailenco 2021-06-11 17:15:06 +03:00
parent 14d82a2d93
commit 652668236d
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
}
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 {
hooks []Hook
hooks []fullHook
}
func (hs *hooks) lock() {
@ -44,7 +58,11 @@ func (hs hooks) clone() hooks {
}
func (hs *hooks) AddHook(hook Hook) {
if hook, ok := hook.(fullHook); ok {
hs.hooks = append(hs.hooks, hook)
} else {
hs.hooks = append(hs.hooks, dummyConnectHook{Hook: hook})
}
}
func (hs hooks) process(
@ -132,6 +150,16 @@ func (hs hooks) withContext(ctx context.Context, fn func() error) error {
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 {