Repository: numaproj/numaflow
Pull request #1040 (2023-09-12 22:03:19)

fix: allow udsource pending api to return negative count to indicate PendingNotAvailable
Such that we distinguish the PendingNotAvailable scenario from there are not pending messages. Hence auto-scaler knows whether to skip scaling or to scale down.

<!--

Before you push your changes:

* Run `make pre-push -B` to fix codegen and lint problems.

Then, you MUST:

* Sign-off your commit (otherwise the DCO check will fail).
* Use [a conventional commit message](https://www.conventionalcommits.org/en/v1.0.0/) (otherwise the commit message check will fail).

If you did not do this, reset all your commit and replace them with a single commit:

```
git reset HEAD~1 ;# change 1 to how many commits you made
git commit --signoff -m 'feat: my feat. Fixes #1234'
```

When creating your PR: 

* Make sure that "Fixes #" or "Closes #" is in both the PR title (for release notes) and description (to automatically link and close the issue).
* Say how you tested your changes. If you changed the UI, attach screenshots.
* Set your PR as a draft initially.
* Your PR needs to pass the required checks before it can be approved. 
* Once required tests have passed, mark your PR "Ready for review".

If changes were requested, once you've made them, you MUST dismiss the review to get it reviewed again.

-->



Patch:
diff --git a/go.mod b/go.mod
index 604f4f91..55a9b7cc 100644
--- a/go.mod
+++ b/go.mod
@@ -27,7 +27,7 @@ require (
 	github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe
 	github.com/nats-io/nats-server/v2 v2.9.19
 	github.com/nats-io/nats.go v1.27.1
-	github.com/numaproj/numaflow-go v0.5.0
+	github.com/numaproj/numaflow-go v0.5.1-0.20230912211616-62600351d97f
 	github.com/prometheus/client_golang v1.14.0
 	github.com/prometheus/common v0.37.0
 	github.com/redis/go-redis/v9 v9.0.3
diff --git a/go.sum b/go.sum
index 82ea5003..e6d65d65 100644
--- a/go.sum
+++ b/go.sum
@@ -670,8 +670,8 @@ github.com/nats-io/nkeys v0.4.4/go.mod h1:XUkxdLPTufzlihbamfzQ7mw/VGx6ObUs+0bN5s
 github.com/nats-io/nuid v1.0.1 h1:5iA8DT8V7q8WK2EScv2padNa/rTESc1KdnPw4TC2paw=
 github.com/nats-io/nuid v1.0.1/go.mod h1:19wcPz3Ph3q0Jbyiqsd0kePYG7A95tJPxeL+1OSON2c=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
-github.com/numaproj/numaflow-go v0.5.0 h1:C49lgMyF7rNdywanzi4NLR+4yxjvDiOGj0IaH/yIF+w=
-github.com/numaproj/numaflow-go v0.5.0/go.mod h1:5zwvvREIbqaCPCKsNE1MVjVToD0kvkCh2Z90Izlhw5U=
+github.com/numaproj/numaflow-go v0.5.1-0.20230912211616-62600351d97f h1:flJEIyki2WF5KuM4sRbvKQwkpEK6TL40lcwNTZqYfbY=
+github.com/numaproj/numaflow-go v0.5.1-0.20230912211616-62600351d97f/go.mod h1:5zwvvREIbqaCPCKsNE1MVjVToD0kvkCh2Z90Izlhw5U=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
 github.com/oklog/ulid v1.3.1 h1:EGfNDEx6MqHz8B3uNV6QAib1UR2Lm97sHi3ocA6ESJ4=
diff --git a/pkg/sources/udsource/grpc_udsource.go b/pkg/sources/udsource/grpc_udsource.go
index 409a3077..e168c88b 100644
--- a/pkg/sources/udsource/grpc_udsource.go
+++ b/pkg/sources/udsource/grpc_udsource.go
@@ -73,7 +73,10 @@ func (u *GRPCBasedUDSource) WaitUntilReady(ctx context.Context) error {
 // ApplyPendingFn returns the number of pending messages in the source.
 func (u *GRPCBasedUDSource) ApplyPendingFn(ctx context.Context) (int64, error) {
 	if resp, err := u.client.PendingFn(ctx, &emptypb.Empty{}); err == nil {
-		return int64(resp.Result.Count), nil
+		if resp.Result.Count < 0 {
+			return isb.PendingNotAvailable, nil
+		}
+		return resp.Result.Count, nil
 	} else {
 		return isb.PendingNotAvailable, err
 	}


Test patch:
diff --git a/pkg/sdkclient/source/client/client_test.go b/pkg/sdkclient/source/client/client_test.go
index bf226a3f..73000129 100644
--- a/pkg/sdkclient/source/client/client_test.go
+++ b/pkg/sdkclient/source/client/client_test.go
@@ -19,8 +19,11 @@ package client
 import (
 	"context"
 	"fmt"
+	"io"
 	"reflect"
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
 	sourcepb "github.com/numaproj/numaflow-go/pkg/apis/proto/source/v1"
@@ -28,6 +31,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/emptypb"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 type rpcMsg struct {
@@ -72,8 +76,70 @@ func TestIsReady(t *testing.T) {
 	assert.EqualError(t, err, "mock connection refused")
 }
 
-// TODO(udsource) - Add test for ReadFn
 func TestReadFn(t *testing.T) {
+	var ctx = context.Background()
+	LintCleanCall()
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClient := sourcemock.NewMockSourceClient(ctrl)
+	mockStreamClient := sourcemock.NewMockSource_ReadFnClient(ctrl)
+
+	var TestEventTime = time.Unix(1661169600, 0).UTC()
+	expectedResp := &sourcepb.ReadResponse{
+		Result: &sourcepb.ReadResponse_Result{
+			Payload:   []byte(`test_payload`),
+			Offset:    &sourcepb.Offset{Offset: []byte(`test_offset`), PartitionId: "0"},
+			EventTime: timestamppb.New(TestEventTime),
+			Keys:      []string{"test_key"},
+		},
+	}
+
+	const numRecords = 5
+	for i := 0; i < numRecords; i++ {
+		mockStreamClient.EXPECT().Recv().Return(expectedResp, nil)
+	}
+	mockStreamClient.EXPECT().Recv().Return(expectedResp, io.EOF)
+
+	mockStreamClient.EXPECT().CloseSend().Return(nil).AnyTimes()
+	mockClient.EXPECT().ReadFn(gomock.Any(), gomock.Any()).Return(mockStreamClient, nil)
+
+	testClient, err := NewFromClient(mockClient)
+	assert.NoError(t, err)
+	assert.True(t, reflect.DeepEqual(testClient, &client{
+		grpcClt: mockClient,
+	}))
+
+	responseCh := make(chan *sourcepb.ReadResponse)
+
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case resp, ok := <-responseCh:
+				if !ok {
+					return
+				}
+				assert.True(t, reflect.DeepEqual(resp, expectedResp))
+			}
+		}
+	}()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		err = testClient.ReadFn(ctx, &sourcepb.ReadRequest{
+			Request: &sourcepb.ReadRequest_Request{
+				NumRecords: uint64(numRecords),
+			},
+		}, responseCh)
+		assert.NoError(t, err)
+	}()
+	wg.Wait()
+	close(responseCh)
 }
 
 func TestAckFn(t *testing.T) {
diff --git a/pkg/sources/udsource/grpc_udsource_test.go b/pkg/sources/udsource/grpc_udsource_test.go
index df711599..22f3ff34 100644
--- a/pkg/sources/udsource/grpc_udsource_test.go
+++ b/pkg/sources/udsource/grpc_udsource_test.go
@@ -87,7 +87,7 @@ func Test_gRPCBasedUDSource_WaitUntilReadyWithMockClient(t *testing.T) {
 }
 
 func Test_gRPCBasedUDSource_ApplyPendingWithMockClient(t *testing.T) {
-	t.Run("test success", func(t *testing.T) {
+	t.Run("test success - positive pending count", func(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		defer ctrl.Finish()
 
@@ -115,6 +115,34 @@ func Test_gRPCBasedUDSource_ApplyPendingWithMockClient(t *testing.T) {
 		assert.Equal(t, int64(123), count)
 	})
 
+	t.Run("test success - pending is not available - negative count", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+		defer ctrl.Finish()
+
+		testResponse := &sourcepb.PendingResponse{
+			Result: &sourcepb.PendingResponse_Result{
+				Count: -123,
+			},
+		}
+
+		mockSourceClient := sourcemock.NewMockSourceClient(ctrl)
+		mockSourceClient.EXPECT().PendingFn(gomock.Any(), gomock.Any()).Return(testResponse, nil).AnyTimes()
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		go func() {
+			<-ctx.Done()
+			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+				t.Log(t.Name(), "test timeout")
+			}
+		}()
+
+		u := NewMockUDSgRPCBasedUDSource(mockSourceClient)
+		count, err := u.ApplyPendingFn(ctx)
+		assert.NoError(t, err)
+		assert.Equal(t, isb.PendingNotAvailable, count)
+	})
+
 	t.Run("test err", func(t *testing.T) {
 		ctrl := gomock.NewController(t)
 		defer ctrl.Finish()