From e1c63729d9d0bcb51a211b505397824826233a35 Mon Sep 17 00:00:00 2001
From: Rain <zhengyu@yunion.cn>
Date: Sun, 26 Apr 2020 20:54:07 +0800
Subject: [PATCH] fix/notify: Add deadline for rpc client

---
 pkg/notify/rpc/apis/send_client.go | 35 ++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/pkg/notify/rpc/apis/send_client.go b/pkg/notify/rpc/apis/send_client.go
index cae3506e87..6391bb1a2d 100644
--- a/pkg/notify/rpc/apis/send_client.go
+++ b/pkg/notify/rpc/apis/send_client.go
@@ -14,16 +14,47 @@
 
 package apis
 
-import "google.golang.org/grpc"
+import (
+	"context"
+	"time"
+
+	"google.golang.org/grpc"
+)
 
 type SendNotificationClient struct {
 	sendAgentClient
-	Conn *grpc.ClientConn
+	Conn        *grpc.ClientConn
+	CallTimeout time.Duration
 }
 
 func NewSendNotificationClient(cc *grpc.ClientConn) *SendNotificationClient {
 	return &SendNotificationClient{
 		sendAgentClient: sendAgentClient{cc},
 		Conn:            cc,
+		CallTimeout:     30 * time.Second,
 	}
 }
+
+func (c *SendNotificationClient) Send(ctx context.Context, in *SendParams, opts ...grpc.CallOption) (*Empty, error) {
+	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	defer cancel()
+	return c.sendAgentClient.Send(ctx, in, opts...)
+}
+
+func (c *SendNotificationClient) UpdateConfig(ctx context.Context, in *UpdateConfigParams, opts ...grpc.CallOption) (*Empty, error) {
+	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	defer cancel()
+	return c.sendAgentClient.UpdateConfig(ctx, in, opts...)
+}
+
+func (c *SendNotificationClient) ValidateConfig(ctx context.Context, in *UpdateConfigParams, opts ...grpc.CallOption) (*ValidateConfigReply, error) {
+	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	defer cancel()
+	return c.sendAgentClient.ValidateConfig(ctx, in, opts...)
+}
+
+func (c *SendNotificationClient) UseridByMobile(ctx context.Context, in *UseridByMobileParams, opts ...grpc.CallOption) (*UseridByMobileReply, error) {
+	ctx, cancel := context.WithTimeout(ctx, c.CallTimeout)
+	defer cancel()
+	return c.sendAgentClient.UseridByMobile(ctx, in, opts...)
+}