-
Notifications
You must be signed in to change notification settings - Fork 53
/
load.go
58 lines (49 loc) · 1.47 KB
/
load.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
package hooks
import (
gsync "github.com/kralicky/gpkg/sync"
"github.com/rancher/opni/pkg/plugins/meta"
"google.golang.org/grpc"
)
type PluginLoadHook interface {
ShouldInvoke(t any) bool
Invoke(t any, md meta.PluginMeta, conn *grpc.ClientConn) (done chan struct{})
}
type onLoadHook[T any] struct {
callback onLoadCallback[T]
}
func (onLoadHook[T]) ShouldInvoke(t any) bool {
_, ok := t.(T)
return ok
}
func (h onLoadHook[T]) Invoke(t any, md meta.PluginMeta, cc *grpc.ClientConn) chan struct{} {
c := make(chan struct{})
go func(t T) {
defer close(c)
h.callback(t, md, cc)
}(t.(T))
return c
}
type onLoadCallback[T any] func(t T, md meta.PluginMeta, conn *grpc.ClientConn)
// Invokes the provided callback function when the plugin of type T is loaded.
func OnLoad[T any](fn func(T)) PluginLoadHook {
return OnLoadMC(func(t T, _ meta.PluginMeta, _ *grpc.ClientConn) {
fn(t)
})
}
// Like OnLoad[T], but adds plugin metadata to the callback.
func OnLoadM[T any](fn func(T, meta.PluginMeta)) PluginLoadHook {
return OnLoadMC(func(t T, md meta.PluginMeta, _ *grpc.ClientConn) {
fn(t, md)
})
}
// Like OnLoadM[T], but adds the grpc client connection to the callback.
func OnLoadMC[T any](fn func(T, meta.PluginMeta, *grpc.ClientConn)) PluginLoadHook {
once := gsync.Map[string, bool]{}
return &onLoadHook[T]{
callback: func(t T, md meta.PluginMeta, cc *grpc.ClientConn) {
if _, loaded := once.LoadOrStore(md.Module, true); !loaded {
fn(t, md, cc)
}
},
}
}