Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 29 additions & 15 deletions pkg/stub/stub.go
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,9 @@ type Stub interface {
// This is the default timeout if the plugin has not been started or
// the timeout received in the Configure request otherwise.
RequestTimeout() time.Duration

// Logger returns the logger used by the stub.
Logger() nrilog.Logger
}

const (
Expand All @@ -188,9 +191,6 @@ const (
)

var (
// Logger for messages generated internally by the stub itself.
log = nrilog.Get()

// Used instead of a nil Context in logging.
noCtx = context.TODO()

Expand Down Expand Up @@ -268,6 +268,14 @@ func WithTTRPCOptions(clientOpts []ttrpc.ClientOpts, serverOpts []ttrpc.ServerOp
}
}

// WithLogger sets the logger to be used by the stub.
func WithLogger(logger nrilog.Logger) Option {
return func(s *stub) error {
s.logger = logger
return nil
}
}

// stub implements Stub.
type stub struct {
sync.Mutex
Expand Down Expand Up @@ -295,6 +303,7 @@ type stub struct {

registrationTimeout time.Duration
requestTimeout time.Duration
logger nrilog.Logger
}

// Handlers for NRI plugin event and request.
Expand Down Expand Up @@ -329,6 +338,7 @@ func New(p interface{}, opts ...Option) (Stub, error) {

registrationTimeout: DefaultRegistrationTimeout,
requestTimeout: DefaultRequestTimeout,
logger: nrilog.Get(),
Copy link
Member

@klihub klihub Oct 15, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gauravkghildiyal I think it would also make sense to add a WithLogger(nrilog.Logger) Option option which would allow setting the used logger with a stub.Option.

}

for _, o := range opts {
Expand All @@ -345,7 +355,7 @@ func New(p interface{}, opts ...Option) (Stub, error) {
return nil, err
}

log.Infof(noCtx, "Created plugin %s (%s, handles %s)", stub.Name(),
stub.logger.Infof(noCtx, "Created plugin %s (%s, handles %s)", stub.Name(),
filepath.Base(os.Args[0]), stub.events.PrettyString())

return stub, nil
Expand Down Expand Up @@ -440,15 +450,15 @@ func (stub *stub) Start(ctx context.Context) (retErr error) {
return err
}

log.Infof(ctx, "Started plugin %s...", stub.Name())
stub.logger.Infof(ctx, "Started plugin %s...", stub.Name())

stub.started = true
return nil
}

// Stop the plugin.
func (stub *stub) Stop() {
log.Infof(noCtx, "Stopping plugin %s...", stub.Name())
stub.logger.Infof(noCtx, "Stopping plugin %s...", stub.Name())

stub.Lock()
defer stub.Unlock()
Expand Down Expand Up @@ -504,7 +514,7 @@ func (stub *stub) Run(ctx context.Context) error {

err = <-stub.srvErrC
if err == ttrpc.ErrServerClosed {
log.Infof(noCtx, "ttrpc server closed %s : %v", stub.Name(), err)
stub.logger.Infof(noCtx, "ttrpc server closed %s : %v", stub.Name(), err)
}

return err
Expand All @@ -522,6 +532,10 @@ func (stub *stub) Name() string {
return stub.idx + "-" + stub.name
}

func (stub *stub) Logger() nrilog.Logger {
return stub.logger
}

func (stub *stub) RegistrationTimeout() time.Duration {
return stub.registrationTimeout
}
Expand All @@ -533,12 +547,12 @@ func (stub *stub) RequestTimeout() time.Duration {
// Connect the plugin to NRI.
func (stub *stub) connect() error {
if stub.conn != nil {
log.Infof(noCtx, "Using given plugin connection...")
stub.logger.Infof(noCtx, "Using given plugin connection...")
return nil
}

if env := os.Getenv(api.PluginSocketEnvVar); env != "" {
log.Infof(noCtx, "Using connection %q from environment...", env)
stub.logger.Infof(noCtx, "Using connection %q from environment...", env)

fd, err := strconv.Atoi(env)
if err != nil {
Expand Down Expand Up @@ -566,7 +580,7 @@ func (stub *stub) connect() error {

// Register the plugin with NRI.
func (stub *stub) register(ctx context.Context) error {
log.Infof(ctx, "Registering plugin %s...", stub.Name())
stub.logger.Infof(ctx, "Registering plugin %s...", stub.Name())

ctx, cancel := context.WithTimeout(ctx, stub.registrationTimeout)
defer cancel()
Expand Down Expand Up @@ -621,7 +635,7 @@ func (stub *stub) Configure(ctx context.Context, req *api.ConfigureRequest) (rpl
err error
)

log.Infof(ctx, "Configuring plugin %s for runtime %s/%s...", stub.Name(),
stub.logger.Infof(ctx, "Configuring plugin %s for runtime %s/%s...", stub.Name(),
req.RuntimeName, req.RuntimeVersion)

stub.registrationTimeout = time.Duration(req.RegistrationTimeout * int64(time.Millisecond))
Expand All @@ -636,7 +650,7 @@ func (stub *stub) Configure(ctx context.Context, req *api.ConfigureRequest) (rpl
} else {
events, err = handler(ctx, req.Config, req.RuntimeName, req.RuntimeVersion)
if err != nil {
log.Errorf(ctx, "Plugin configuration failed: %v", err)
stub.logger.Errorf(ctx, "Plugin configuration failed: %v", err)
return nil, err
}

Expand All @@ -646,13 +660,13 @@ func (stub *stub) Configure(ctx context.Context, req *api.ConfigureRequest) (rpl

// Only allow plugins to subscribe to events they can handle.
if extra := events & ^stub.events; extra != 0 {
log.Errorf(ctx, "Plugin subscribed for unhandled events %s (0x%x)",
stub.logger.Errorf(ctx, "Plugin subscribed for unhandled events %s (0x%x)",
extra.PrettyString(), extra)
return nil, fmt.Errorf("internal error: unhandled events %s (0x%x)",
extra.PrettyString(), extra)
}

log.Infof(ctx, "Subscribing plugin %s (%s) for events %s", stub.Name(),
stub.logger.Infof(ctx, "Subscribing plugin %s (%s) for events %s", stub.Name(),
filepath.Base(os.Args[0]), events.PrettyString())
}

Expand All @@ -679,7 +693,7 @@ func (stub *stub) collectSync(req *api.SynchronizeRequest) (*api.SynchronizeResp
stub.Lock()
defer stub.Unlock()

log.Debugf(noCtx, "collecting sync req with %d pods, %d containers...",
stub.logger.Debugf(noCtx, "collecting sync req with %d pods, %d containers...",
len(req.Pods), len(req.Containers))

if stub.syncReq == nil {
Expand Down
Loading