From dd0e574a6aa416db2fdfcbfc6f8affac694c2a07 Mon Sep 17 00:00:00 2001
From: wulixuan <931845457@qq.com>
Date: Tue, 7 Jun 2022 11:48:39 +0800
Subject: [PATCH] add grpc timeout

---
 cmd/client/client.go | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/cmd/client/client.go b/cmd/client/client.go
index 69b57792..7fea261b 100644
--- a/cmd/client/client.go
+++ b/cmd/client/client.go
@@ -20,6 +20,7 @@ package main
 import (
 	"context"
 	goflag "flag"
+	"time"
 
 	vcudaapi "tkestack.io/gpu-manager/pkg/api/runtime/vcuda"
 	"tkestack.io/gpu-manager/pkg/flags"
@@ -35,6 +36,8 @@ var (
 	addr, busID, podUID, contName, contID string
 )
 
+const TimeOut = 1 * time.Second
+
 func main() {
 	cmdFlags := pflag.CommandLine
 
@@ -53,14 +56,14 @@ func main() {
 		klog.Fatalf("argument is empty, current: %s", cmdFlags.Args())
 	}
 
-	conn, err := grpc.Dial(addr, utils.DefaultDialOptions...)
+	options := append(utils.DefaultDialOptions, grpc.WithTimeout(TimeOut))
+	conn, err := grpc.Dial(addr, options...)
 	if err != nil {
 		klog.Fatalf("can't dial %s, error %v", addr, err)
 	}
 	defer conn.Close()
 
 	client := vcudaapi.NewVCUDAServiceClient(conn)
-	ctx := context.TODO()
 
 	req := &vcudaapi.VDeviceRequest{
 		BusId:         busID,
@@ -73,6 +76,8 @@ func main() {
 		req.ContainerId = contID
 	}
 
+	ctx, cancel := context.WithTimeout(context.Background(), TimeOut)
+	defer cancel()
 	_, err = client.RegisterVDevice(ctx, req)
 	if err != nil {
 		klog.Fatalf("fail to get response from manager, error %v", err)