Add timeout for discoverNetworkServiceEndpoint (#986)

Signed-off-by: Vladimir Popov <vladimir.popov@xored.com>

pkg/networkservice/common/discover/server.go

@@ -127,6 +127,8 @@ func (d *discoverCandidatesServer) Close(ctx context.Context, conn *networkservi
 }
 
 func (d *discoverCandidatesServer) discoverNetworkServiceEndpoint(ctx context.Context, nseName string) (*registry.NetworkServiceEndpoint, error) {
+	clockTime := clock.FromContext(ctx)
+
 	query := &registry.NetworkServiceEndpointQuery{
 		NetworkServiceEndpoint: &registry.NetworkServiceEndpoint{
 			Name: nseName,
@@ -147,7 +149,17 @@ func (d *discoverCandidatesServer) discoverNetworkServiceEndpoint(ctx context.Co
 
 	query.Watch = true
 
-	nseStream, err = d.nseClient.Find(ctx, query)
+	findTimeout := 100 * time.Millisecond
+	if deadline, ok := ctx.Deadline(); ok {
+		if ctxTimeout := clockTime.Until(deadline) / 10; ctxTimeout > findTimeout {
+			findTimeout = ctxTimeout
+		}
+	}
+
+	findCtx, cancelFind := clock.FromContext(ctx).WithTimeout(ctx, findTimeout)
+	defer cancelFind()
+
+	nseStream, err = d.nseClient.Find(findCtx, query)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

pkg/networkservice/common/discover/server_test.go

@@ -474,6 +474,9 @@ func TestMatchExactService(t *testing.T) {
 func TestMatchExactEndpoint(t *testing.T) {
 	t.Cleanup(func() { goleak.VerifyNone(t) })
 
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
 	nseServer := memory.NewNetworkServiceEndpointRegistryServer()
 
 	nseName := "final-endpoint"
@@ -497,9 +500,6 @@ func TestMatchExactEndpoint(t *testing.T) {
 	require.NoError(t, err)
 
 	// 2. Try to discover NSE by the right name
-	ctx, cancel := context.WithTimeout(context.Background(), testWait)
-	defer cancel()
-
 	request := &networkservice.NetworkServiceRequest{
 		Connection: &networkservice.Connection{
 			NetworkServiceEndpointName: nseName,
@@ -517,9 +517,6 @@ func TestMatchExactEndpoint(t *testing.T) {
 	require.NoError(t, err)
 
 	// 4. Try to discover NSE by the right name
-	ctx, cancel = context.WithTimeout(context.Background(), testWait)
-	defer cancel()
-
 	_, err = server.Request(ctx, request.Clone())
 	require.NoError(t, err)
 }