pkg/adaptation: add support for extra ttrpc options.

klihub · klihub · commit c4e2f81f94b3 · 2024-01-30T14:20:48.000+02:00
Add WithTTRPCOptions() which can be used to pass extra client
and server options to ttrpc.

Signed-off-by: Krisztian Litkey &lt;krisztian.litkey@intel.com&gt;
diff --git a/pkg/adaptation/adaptation.go b/pkg/adaptation/adaptation.go
@@ -29,6 +29,7 @@ import (
 
 	"github.com/containerd/nri/pkg/api"
 	"github.com/containerd/nri/pkg/log"
+	"github.com/containerd/ttrpc"
 )
 
 const (
@@ -60,6 +61,8 @@ type Adaptation struct {
 	dontListen bool
 	syncFn     SyncFn
 	updateFn   UpdateFn
+	clientOpts []ttrpc.ClientOpts
+	serverOpts []ttrpc.ServerOpt
 	listener   net.Listener
 	plugins    []*plugin
 }
@@ -104,6 +107,15 @@ func WithDisabledExternalConnections() Option {
 	}
 }
 
+// WithTTRPCOptions sets extra client and server options to use for ttrpc.
+func WithTTRPCOptions(clientOpts []ttrpc.ClientOpts, serverOpts []ttrpc.ServerOpt) Option {
+	return func(r *Adaptation) error {
+		r.clientOpts = append(r.clientOpts, clientOpts...)
+		r.serverOpts = append(r.serverOpts, serverOpts...)
+		return nil
+	}
+}
+
 // New creates a new NRI Runtime.
 func New(name, version string, syncFn SyncFn, updateFn UpdateFn, opts ...Option) (*Adaptation, error) {
 	var err error
diff --git a/pkg/adaptation/plugin.go b/pkg/adaptation/plugin.go
@@ -198,20 +198,24 @@ func (p *plugin) connect(conn stdnet.Conn) (retErr error) {
 	if err != nil {
 		return fmt.Errorf("failed to mux plugin connection for plugin %q: %w", p.name(), err)
 	}
-	rpcc := ttrpc.NewClient(pconn, ttrpc.WithOnClose(
-		func() {
-			log.Infof(noCtx, "connection to plugin %q closed", p.name())
-			close(p.closeC)
-			p.close()
-		}))
+
+	clientOpts := []ttrpc.ClientOpts{
+		ttrpc.WithOnClose(
+			func() {
+				log.Infof(noCtx, "connection to plugin %q closed", p.name())
+				close(p.closeC)
+				p.close()
+			}),
+	}
+	rpcc := ttrpc.NewClient(pconn, append(clientOpts, p.r.clientOpts...)...)
 	defer func() {
 		if retErr != nil {
 			rpcc.Close()
 		}
 	}()
 	stub := api.NewPluginClient(rpcc)
 
-	rpcs, err := ttrpc.NewServer()
+	rpcs, err := ttrpc.NewServer(p.r.serverOpts...)
 	if err != nil {
 		return fmt.Errorf("failed to create ttrpc server for plugin %q: %w", p.name(), err)
 	}
