@@ -17,8 +17,7 @@ import (
1717
1818 securejoin "github.com/cyphar/filepath-securejoin"
1919 "github.com/onkernel/hypeman/lib/hypervisor"
20- "google.golang.org/grpc"
21- "google.golang.org/grpc/credentials/insecure"
20+ "storj.io/drpc/drpcconn"
2221)
2322
2423const (
@@ -51,25 +50,31 @@ func IsAgentConnectionError(err error) bool {
5150 errors .As (err , & agentErr ))
5251}
5352
54- // connPool manages reusable gRPC connections per vsock dialer key
53+ // pooledConn wraps a drpc connection with its underlying net.Conn for cleanup
54+ type pooledConn struct {
55+ drpc * drpcconn.Conn
56+ netConn net.Conn
57+ }
58+
59+ // connPool manages reusable DRPC connections per vsock dialer key
5560// This avoids the overhead and potential issues of rapidly creating/closing connections
5661var connPool = struct {
5762 sync.RWMutex
58- conns map [string ]* grpc. ClientConn
63+ conns map [string ]* pooledConn
5964}{
60- conns : make (map [string ]* grpc. ClientConn ),
65+ conns : make (map [string ]* pooledConn ),
6166}
6267
6368// GetOrCreateConn returns an existing connection or creates a new one using a VsockDialer.
6469// This supports multiple hypervisor types (Cloud Hypervisor, QEMU, etc.).
65- func GetOrCreateConn (ctx context.Context , dialer hypervisor.VsockDialer ) (* grpc. ClientConn , error ) {
70+ func GetOrCreateConn (ctx context.Context , dialer hypervisor.VsockDialer ) (* drpcconn. Conn , error ) {
6671 key := dialer .Key ()
6772
6873 // Try read lock first for existing connection
6974 connPool .RLock ()
70- if conn , ok := connPool .conns [key ]; ok {
75+ if pc , ok := connPool .conns [key ]; ok {
7176 connPool .RUnlock ()
72- return conn , nil
77+ return pc . drpc , nil
7378 }
7479 connPool .RUnlock ()
7580
@@ -78,40 +83,33 @@ func GetOrCreateConn(ctx context.Context, dialer hypervisor.VsockDialer) (*grpc.
7883 defer connPool .Unlock ()
7984
8085 // Double-check after acquiring write lock
81- if conn , ok := connPool .conns [key ]; ok {
82- return conn , nil
86+ if pc , ok := connPool .conns [key ]; ok {
87+ return pc . drpc , nil
8388 }
8489
8590 // Create new connection using the VsockDialer
86- conn , err := grpc .Dial ("passthrough:///vsock" ,
87- grpc .WithContextDialer (func (ctx context.Context , addr string ) (net.Conn , error ) {
88- netConn , err := dialer .DialVsock (ctx , vsockGuestPort )
89- if err != nil {
90- return nil , & AgentConnectionError {Err : err }
91- }
92- return netConn , nil
93- }),
94- grpc .WithTransportCredentials (insecure .NewCredentials ()),
95- )
91+ netConn , err := dialer .DialVsock (ctx , vsockGuestPort )
9692 if err != nil {
97- return nil , fmt . Errorf ( "create grpc connection: %w" , err )
93+ return nil , & AgentConnectionError { Err : err }
9894 }
9995
100- connPool .conns [key ] = conn
101- slog .Debug ("created new gRPC connection" , "key" , key )
96+ // Wrap in drpc connection
97+ conn := drpcconn .New (netConn )
98+
99+ connPool .conns [key ] = & pooledConn {drpc : conn , netConn : netConn }
100+ slog .Debug ("created new DRPC connection" , "key" , key )
102101 return conn , nil
103102}
104103
105104// CloseConn removes a connection from the pool by key (call when VM is deleted).
106- // We only remove from pool, not explicitly close - the connection will fail
107- // naturally when the VM dies, and grpc will clean up.
108105func CloseConn (dialerKey string ) {
109106 connPool .Lock ()
110107 defer connPool .Unlock ()
111108
112- if _ , ok := connPool .conns [dialerKey ]; ok {
109+ if pc , ok := connPool .conns [dialerKey ]; ok {
110+ pc .drpc .Close ()
113111 delete (connPool .conns , dialerKey )
114- slog .Debug ("removed gRPC connection from pool" , "key" , dialerKey )
112+ slog .Debug ("removed DRPC connection from pool" , "key" , dialerKey )
115113 }
116114}
117115
@@ -132,27 +130,27 @@ type ExecOptions struct {
132130 Timeout int32 // Execution timeout in seconds (0 = no timeout)
133131}
134132
135- // ExecIntoInstance executes command in instance via vsock using gRPC .
133+ // ExecIntoInstance executes command in instance via vsock using DRPC .
136134// The dialer is a hypervisor-specific VsockDialer that knows how to connect to the guest.
137135func ExecIntoInstance (ctx context.Context , dialer hypervisor.VsockDialer , opts ExecOptions ) (* ExitStatus , error ) {
138136 start := time .Now ()
139137 var bytesSent int64
140138
141- // Get or create a reusable gRPC connection for this vsock dialer
142- grpcConn , err := GetOrCreateConn (ctx , dialer )
139+ // Get or create a reusable DRPC connection for this vsock dialer
140+ drpcConn , err := GetOrCreateConn (ctx , dialer )
143141 if err != nil {
144- return nil , fmt .Errorf ("get grpc connection: %w" , err )
142+ return nil , fmt .Errorf ("get drpc connection: %w" , err )
145143 }
146144 // Note: Don't close the connection - it's pooled and reused
147145
148146 // Create guest client
149- client := NewGuestServiceClient ( grpcConn )
147+ client := NewDRPCGuestServiceClient ( drpcConn )
150148 stream , err := client .Exec (ctx )
151149 if err != nil {
152150 return nil , fmt .Errorf ("start exec stream: %w" , err )
153151 }
154152 // Ensure stream is properly closed when we're done
155- defer stream .CloseSend ()
153+ defer stream .Close ()
156154
157155 // Send start request
158156 if err := stream .Send (& ExecRequest {
@@ -233,12 +231,12 @@ type CopyToInstanceOptions struct {
233231// CopyToInstance copies a file or directory to an instance via vsock.
234232// The dialer is a hypervisor-specific VsockDialer that knows how to connect to the guest.
235233func CopyToInstance (ctx context.Context , dialer hypervisor.VsockDialer , opts CopyToInstanceOptions ) error {
236- grpcConn , err := GetOrCreateConn (ctx , dialer )
234+ drpcConn , err := GetOrCreateConn (ctx , dialer )
237235 if err != nil {
238- return fmt .Errorf ("get grpc connection: %w" , err )
236+ return fmt .Errorf ("get drpc connection: %w" , err )
239237 }
240238
241- client := NewGuestServiceClient ( grpcConn )
239+ client := NewDRPCGuestServiceClient ( drpcConn )
242240
243241 // Stat the source
244242 srcInfo , err := os .Stat (opts .SrcPath )
@@ -253,7 +251,7 @@ func CopyToInstance(ctx context.Context, dialer hypervisor.VsockDialer, opts Cop
253251}
254252
255253// copyFileToInstance copies a single file to the instance
256- func copyFileToInstance (ctx context.Context , client GuestServiceClient , srcPath , dstPath string , mode fs.FileMode ) error {
254+ func copyFileToInstance (ctx context.Context , client DRPCGuestServiceClient , srcPath , dstPath string , mode fs.FileMode ) error {
257255 srcInfo , err := os .Stat (srcPath )
258256 if err != nil {
259257 return fmt .Errorf ("stat source: %w" , err )
@@ -329,7 +327,7 @@ func copyFileToInstance(ctx context.Context, client GuestServiceClient, srcPath,
329327}
330328
331329// copyDirToInstance copies a directory recursively to the instance
332- func copyDirToInstance (ctx context.Context , client GuestServiceClient , srcPath , dstPath string ) error {
330+ func copyDirToInstance (ctx context.Context , client DRPCGuestServiceClient , srcPath , dstPath string ) error {
333331 srcPath = filepath .Clean (srcPath )
334332
335333 // First create the destination directory
@@ -444,12 +442,12 @@ type FileHandler func(header *CopyFromGuestHeader, data io.Reader) error
444442// CopyFromInstance copies a file or directory from an instance via vsock.
445443// The dialer is a hypervisor-specific VsockDialer that knows how to connect to the guest.
446444func CopyFromInstance (ctx context.Context , dialer hypervisor.VsockDialer , opts CopyFromInstanceOptions ) error {
447- grpcConn , err := GetOrCreateConn (ctx , dialer )
445+ drpcConn , err := GetOrCreateConn (ctx , dialer )
448446 if err != nil {
449- return fmt .Errorf ("get grpc connection: %w" , err )
447+ return fmt .Errorf ("get drpc connection: %w" , err )
450448 }
451449
452- client := NewGuestServiceClient ( grpcConn )
450+ client := NewDRPCGuestServiceClient ( drpcConn )
453451
454452 stream , err := client .CopyFromGuest (ctx , & CopyFromGuestRequest {
455453 Path : opts .SrcPath ,
0 commit comments