@@ -19,7 +19,9 @@ package quota_test
1919import (
2020 "strings"
2121 "testing"
22+ "time"
2223
24+ "github.com/eapache/go-resiliency/retrier"
2325 "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/quotaplugins/quota-forest/quota-manager/quota"
2426 "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/quotaplugins/quota-forest/quota-manager/quota/utils"
2527 "github.com/stretchr/testify/assert"
@@ -225,3 +227,169 @@ func TestNewQuotaManagerConsumerAllocationRelease(t *testing.T) {
225227 })
226228 }
227229}
230+ func TestQuotaManagerQuotaUsedLongRunningConsumers (t * testing.T ) {
231+ forestName := "unit-test-1"
232+ qmManagerUnderTest := quota .NewManager ()
233+
234+ err := qmManagerUnderTest .AddForest (forestName )
235+ assert .NoError (t , err , "No error expected when adding a forest" )
236+ testTreeName , err := qmManagerUnderTest .AddTreeFromString (
237+ `{
238+ "kind": "QuotaTree",
239+ "metadata": {
240+ "name": "test-tree"
241+ },
242+ "spec": {
243+ "resourceNames": [
244+ "cpu",
245+ "memory"
246+ ],
247+ "nodes": {
248+ "root": {
249+ "parent": "nil",
250+ "hard": "true",
251+ "quota": {
252+ "cpu": "10",
253+ "memory": "256"
254+ }
255+ },
256+ "gold": {
257+ "parent": "root",
258+ "hard": "true",
259+ "quota": {
260+ "cpu": "10",
261+ "memory": "256"
262+ }
263+ }
264+ }
265+ }
266+ }` )
267+ if err != nil {
268+ assert .Fail (t , "No error expected when adding a tree to forest" )
269+ }
270+ err = qmManagerUnderTest .AddTreeToForest (forestName , testTreeName )
271+ assert .NoError (t , err , "No error expected when adding a tree from forest" )
272+ modeSet := qmManagerUnderTest .SetMode (quota .Normal )
273+ assert .True (t , modeSet , "Setting the mode should not fail." )
274+
275+ // Define the test table
276+ var tests = []struct {
277+ name string
278+ consumer utils.JConsumer
279+ }{
280+ // the table itself
281+ {"Gold consumer 1" ,
282+ utils.JConsumer {
283+ Kind : "Consumer" ,
284+ MetaData : utils.JMetaData {
285+ Name : "gpld-consumer-data" ,
286+ },
287+ Spec : utils.JConsumerSpec {
288+ ID : "gold-consumer-1" ,
289+ Trees : []utils.JConsumerTreeSpec {
290+ {
291+ TreeName : testTreeName ,
292+ GroupID : "gold" ,
293+ Request : map [string ]int {
294+ "cpu" : 10 ,
295+ "memory" : 4 ,
296+ "gpu" : 0 ,
297+ },
298+ Priority : 0 ,
299+ CType : 0 ,
300+ UnPreemptable : false ,
301+ },
302+ },
303+ },
304+ },
305+ },
306+ // the table itself
307+ {"Gold consumer 2" ,
308+ utils.JConsumer {
309+ Kind : "Consumer" ,
310+ MetaData : utils.JMetaData {
311+ Name : "gpld-consumer-data" ,
312+ },
313+ Spec : utils.JConsumerSpec {
314+ ID : "gold-consumer-2" ,
315+ Trees : []utils.JConsumerTreeSpec {
316+ {
317+ TreeName : testTreeName ,
318+ GroupID : "gold" ,
319+ Request : map [string ]int {
320+ "cpu" : 10 ,
321+ "memory" : 4 ,
322+ "gpu" : 0 ,
323+ },
324+ Priority : 0 ,
325+ CType : 0 ,
326+ UnPreemptable : false ,
327+ },
328+ },
329+ },
330+ },
331+ },
332+ }
333+ // Execute tests in parallel
334+ for _ , tc := range tests {
335+ tc := tc // capture range variable
336+ t .Run (tc .name , func (t * testing.T ) {
337+ t .Parallel ()
338+ // Get list of quota management tree IDs
339+ qmTreeIDs := qmManagerUnderTest .GetTreeNames ()
340+
341+ consumerInfo , err := quota .NewConsumerInfo (tc .consumer )
342+ assert .NoError (t , err , "No error expected when building consumer" )
343+ assert .Contains (t , qmTreeIDs , tc .consumer .Spec .Trees [0 ].TreeName )
344+
345+ retryAllocation := retrier .New (retrier .LimitedExponentialBackoff (10 , 10 * time .Millisecond , 1 * time .Second ),
346+ & AllocationClassifier {})
347+
348+ err = retryAllocation .Run (func () error {
349+ added , err2 := qmManagerUnderTest .AddConsumer (consumerInfo )
350+ assert .NoError (t , err2 , "No error expected when adding consumer" )
351+ assert .True (t , added , "Consumer is expected to be added" )
352+
353+ response , err2 := qmManagerUnderTest .AllocateForest (forestName , consumerInfo .GetID ())
354+ if err2 == nil {
355+ assert .Equal (t , 0 , len (strings .TrimSpace (response .GetMessage ())), "A empty response is expected" )
356+ assert .True (t , response .IsAllocated (), "The allocation should succeed" )
357+ } else {
358+ removed , err3 := qmManagerUnderTest .RemoveConsumer (consumerInfo .GetID ())
359+ assert .NoError (t , err3 , "No Error expected when removing consumer" )
360+ assert .True (t , removed , "Removal of consumer should succeed" )
361+ }
362+ return err2
363+
364+ })
365+ if err != nil {
366+ assert .Failf (t , "Allocation of quota should have succeed: '%s'" , err .Error ())
367+ }
368+
369+ //simulate a long running consumer that has quota allocated
370+ time .Sleep (10 * time .Millisecond )
371+
372+ deAllocated := qmManagerUnderTest .DeAllocateForest (forestName , consumerInfo .GetID ())
373+ assert .True (t , deAllocated , "De-allocation expected to succeed" )
374+
375+ removed , err := qmManagerUnderTest .RemoveConsumer (consumerInfo .GetID ())
376+ assert .NoError (t , err , "No Error expected when removing consumer" )
377+ assert .True (t , removed , "Removal of consumer should succeed" )
378+
379+ })
380+ }
381+
382+ }
383+
384+ type AllocationClassifier struct {
385+ }
386+
387+ func (c * AllocationClassifier ) Classify (err error ) retrier.Action {
388+ if err == nil {
389+ return retrier .Succeed
390+ }
391+ if strings .TrimSpace (err .Error ()) == "Failed to allocate quota on quota designation 'test-tree'" {
392+ return retrier .Retry
393+ }
394+ return retrier .Fail
395+ }
0 commit comments