Create RSA keys on VizierClient creation instead of every ExecuteScript

Summary: TSIA, creating RSA keys is expensive, let's cache them. Test Plan: Run the go api examples Reviewers: zasgar, philkuz Reviewed By: zasgar Signed-off-by: Vihang Mehta <[email protected]> Differential Revision: https://phab.corp.pixielabs.ai/D9908 GitOrigin-RevId: 5366ed29ad0fab18468cd3217e3abfdf6e55797d
pixie-io · Oct 5, 2021 · b159812 · b159812
1 parent 4571dff
commit b159812
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 19 deletions.
diff --git a/client.go b/client.go
@@ -30,6 +30,7 @@ import (
 	"google.golang.org/grpc/metadata"
 
 	"px.dev/pxapi/types"
+	"px.dev/pxapi/utils"
 	"px.dev/pxapi/proto/cloudpb"
 	"px.dev/pxapi/proto/vizierpb"
 )
@@ -148,12 +149,21 @@ func (c *Client) NewVizierClient(ctx context.Context, vizierID string) (*VizierC
 		vzConn = conn
 	}
 
+	var encOpts, decOpts *vizierpb.ExecuteScriptRequest_EncryptionOptions
+	if c.useEncryption {
+		encOpts, decOpts, err = utils.CreateEncryptionOptions()
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	// Now create the actual client.
 	vzClient := &VizierClient{
-		cloud:         c,
-		useEncryption: c.useEncryption,
-		vizierID:      vizierID,
-		vzClient:      vizierpb.NewVizierServiceClient(vzConn),
+		cloud:    c,
+		encOpts:  encOpts,
+		decOpts:  decOpts,
+		vizierID: vizierID,
+		vzClient: vizierpb.NewVizierServiceClient(vzConn),
 	}
 
 	return vzClient, nil

diff --git a/vizier.go b/vizier.go
@@ -21,32 +21,24 @@ package pxapi
 import (
 	"context"
 
-	"px.dev/pxapi/utils"
 	"px.dev/pxapi/proto/vizierpb"
 )
 
 // VizierClient is the client for a single vizier.
 type VizierClient struct {
-	cloud         *Client
-	useEncryption bool
-	vizierID      string
-	vzClient      vizierpb.VizierServiceClient
+	cloud    *Client
+	vizierID string
+	vzClient vizierpb.VizierServiceClient
+	encOpts  *vizierpb.ExecuteScriptRequest_EncryptionOptions
+	decOpts  *vizierpb.ExecuteScriptRequest_EncryptionOptions
 }
 
 // ExecuteScript runs the script on vizier.
 func (v *VizierClient) ExecuteScript(ctx context.Context, pxl string, mux TableMuxer) (*ScriptResults, error) {
-	var encOpts, decOpts *vizierpb.ExecuteScriptRequest_EncryptionOptions
-	var err error
-	if v.useEncryption {
-		encOpts, decOpts, err = utils.CreateEncryptionOptions()
-		if err != nil {
-			return nil, err
-		}
-	}
 	req := &vizierpb.ExecuteScriptRequest{
 		ClusterID:         v.vizierID,
 		QueryStr:          pxl,
-		EncryptionOptions: encOpts,
+		EncryptionOptions: v.encOpts,
 	}
 	ctx, cancel := context.WithCancel(ctx)
 	res, err := v.vzClient.ExecuteScript(v.cloud.cloudCtxWithMD(ctx), req)
@@ -59,7 +51,7 @@ func (v *VizierClient) ExecuteScript(ctx context.Context, pxl string, mux TableM
 	sr.c = res
 	sr.cancel = cancel
 	sr.tm = mux
-	sr.decOpts = decOpts
+	sr.decOpts = v.decOpts
 
 	return sr, nil
 }