Repository: armadaproject/armada

Issue:
Armada Scheduler Slow Memory Leak
I've observed that the Armada Scheduler has a slow memory leak. I just enabled pprof so I should have some results in the next few days but wanted to open this for tracking. This is on v0.16.9 but we've been seeing in in older versions as well.

![Image](https://github.com/user-attachments/assets/605e4d59-0073-418f-9eea-cba38c2875fb)

Hints from issue comments:
@Sovietaced: Turns out this memory leak is caused by a leak in the regular calls to fetch the latest queues.

![Image](https://github.com/user-attachments/assets/1abee326-471a-483b-9dde-ef19cee5c596)@dejanzele: Hey,

Thanks for raising this PR!

Somebody from the team will check this out and review.

Solution patch:
diff --git a/internal/scheduler/queue/queue_cache.go b/internal/scheduler/queue/queue_cache.go
index 9a07a5af..65d1ea9f 100644
--- a/internal/scheduler/queue/queue_cache.go
+++ b/internal/scheduler/queue/queue_cache.go
@@ -2,6 +2,7 @@ package queue
 
 import (
 	"fmt"
+	"io"
 	"sync/atomic"
 	"time"
 
@@ -72,10 +73,12 @@ func (c *ApiQueueCache) fetchQueues(ctx *armadacontext.Context) error {
 		return err
 	}
 	queues := make([]*api.Queue, 0)
-	for {
-		msg, err := stream.Recv()
-		if err != nil {
-			return err
+	shouldEndStream := false
+
+	for !shouldEndStream {
+		msg, recvErr := stream.Recv()
+		if recvErr != nil {
+			return recvErr
 		}
 		switch msg.GetEvent().(type) {
 		case *api.StreamingQueueMessage_Queue:
@@ -83,9 +86,31 @@ func (c *ApiQueueCache) fetchQueues(ctx *armadacontext.Context) error {
 		case *api.StreamingQueueMessage_End:
 			c.queues.Store(&queues)
 			ctx.Infof("Refreshed Queues in %s", time.Since(start))
-			return nil
+			shouldEndStream = true
 		default:
 			return fmt.Errorf("unknown event of type %T", msg.GetEvent())
 		}
 	}
+
+	if closeErr := closeStream(stream); closeErr != nil {
+		return fmt.Errorf("closing get queues stream: %w", closeErr)
+	}
+
+	return nil
+}
+
+// This should be called after our end of stream message has been seen (StreamingQueueMessage_End)
+// We call recv one more time and expect an EOF back, indicating the stream is properly closed
+func closeStream(stream api.Submit_GetQueuesClient) error {
+	res, err := stream.Recv()
+	if err == nil {
+		switch typed := res.GetEvent().(type) {
+		default:
+			return fmt.Errorf("failed closing stream - unexpectedly received event of type %T", typed)
+		}
+	} else if err == io.EOF {
+		return nil
+	} else {
+		return err
+	}
 }



Test patch:
diff --git a/internal/scheduler/queue/queue_cache_test.go b/internal/scheduler/queue/queue_cache_test.go
index d57c65af..63218b4d 100644
--- a/internal/scheduler/queue/queue_cache_test.go
+++ b/internal/scheduler/queue/queue_cache_test.go
@@ -2,6 +2,7 @@ package queue
 
 import (
 	"fmt"
+	"io"
 	"testing"
 	"time"
 
@@ -61,13 +62,22 @@ func TestFetch(t *testing.T) {
 					Recv().
 					Return(nil, fmt.Errorf("dummy error"))
 			} else {
-				mockStream.
+				appEnd := mockStream.
 					EXPECT().
 					Recv().
 					Return(
 						&api.StreamingQueueMessage{
 							Event: &api.StreamingQueueMessage_End{},
 						}, nil)
+				grpcEnd := mockStream.
+					EXPECT().
+					Recv().
+					Return(nil, io.EOF)
+
+				gomock.InOrder(
+					appEnd,
+					grpcEnd,
+				)
 			}
 
 			mockApiClient.EXPECT().GetQueues(ctx, gomock.Any()).Return(mockStream, nil).Times(1)