@@ -90,7 +90,7 @@ func TestS3GetAllConcurrently(t *testing.T) {
9090 }
9191
9292 // ASSERT input and output order is the same.
93- require .Equal (t , len (outputKeys ), total )
93+ require .Equal (t , total , len (outputKeys ))
9494 for i := 0 ; i < total ; i ++ {
9595 assert .Equal (t , aws .ToString (objects [i ].Key ), outputKeys [i ])
9696 }
@@ -123,6 +123,117 @@ func TestS3GetAllConcurrently(t *testing.T) {
123123 }
124124}
125125
126+ // go test --run TestS3GetAllConcurrentlyWithContext -v
127+ func TestS3GetAllConcurrentlyWithContext (t * testing.T ) {
128+ // ARRANGE
129+ setup ()
130+ defer teardown ()
131+
132+ // ASSERT parameter errors.
133+ _ , err := NewConcurrent (0 , 100 , 1000 )
134+ assert .NotNil (t , err )
135+ _ , err = NewConcurrent (100 , 0 , 1000 )
136+ assert .NotNil (t , err )
137+ _ , err = NewConcurrent (100 , 100 , 0 )
138+ assert .NotNil (t , err )
139+ _ , err = NewConcurrent (100 , 10 , 99 )
140+ assert .NotNil (t , err )
141+ _ , err = NewConcurrent (100 , 101 , 1000 )
142+ assert .NotNil (t , err )
143+
144+ client , err := NewConcurrent (100 , 10 , 1000 )
145+ require .Nil (t , err , fmt .Sprintf ("error creating s3 client concurrency manager: %v" , err ))
146+
147+ // ASSERT computed fields.
148+ assert .Equal (t , 100 , len (client .manager .workerPool .channel ))
149+ assert .Equal (t , 100 , len (client .manager .memoryPool .channel ))
150+ assert .Equal (t , int64 (10 ), client .manager .memoryChunkSize )
151+ assert .Equal (t , int64 (10 * 100 ), client .manager .memoryTotalSize )
152+ assert .Equal (t , 10 , client .manager .maxWorkersPerRequest )
153+
154+ // ASSERT memory chunk size is correct in memory pool.
155+ chunk := <- client .manager .memoryPool .channel
156+ assert .Equal (t , int64 (10 ), chunk )
157+ client .manager .memoryPool .channel <- chunk
158+
159+ // ASSERT worker get/release methods work expectedly.
160+ w := client .manager .getWorkers (1 )
161+ assert .Equal (t , 99 , len (client .manager .workerPool .channel ))
162+ client .manager .returnWorker (w [0 ])
163+ assert .Equal (t , 100 , len (client .manager .workerPool .channel ))
164+
165+ // ASSERT memory get/release methods work expectedly.
166+ elevenByteFile := types.Object {Size : aws .Int64 (11 )} // requires 2 memory chunks.
167+ client .manager .secureMemory ([]types.Object {elevenByteFile })
168+ assert .Equal (t , 98 , len (client .manager .memoryPool .channel ))
169+ client .manager .releaseMemory (20 )
170+ assert .Equal (t , 100 , len (client .manager .memoryPool .channel ))
171+
172+ // ARRANGE bucket with test objects.
173+ total := 20
174+ keys := make ([]string , total )
175+ for i := 0 ; i < total ; i ++ {
176+ keys [i ] = fmt .Sprintf ("%s-%v" , testObjectKey , i )
177+ }
178+ awsCmdPutKeys (keys )
179+
180+ // ACTION
181+ objects , _ := client .ListAllObjects (testBucket , "" )
182+ tooManyBytes := make ([]types.Object , 10 * len (objects ))
183+ for _ , o := range objects {
184+ for i := 0 ; i < 10 ; i ++ {
185+ tooManyBytes = append (tooManyBytes , o )
186+ }
187+ }
188+ output := client .GetAllConcurrentlyWithContext (context .Background (), testBucket , "" , tooManyBytes )
189+
190+ // ASSERT error returned
191+ for hf := range output {
192+ assert .NotNil (t , hf .Error )
193+ }
194+
195+ // ACTION
196+ objects , _ = client .ListAllObjects (testBucket , "" )
197+ output = client .GetAllConcurrentlyWithContext (context .Background (), testBucket , "" , objects )
198+ outputKeys := make ([]string , 0 )
199+ for hf := range output {
200+ outputKeys = append (outputKeys , hf .Key )
201+ }
202+
203+ // ASSERT input and output order is the same.
204+ require .Equal (t , total , len (outputKeys ))
205+ for i := 0 ; i < total ; i ++ {
206+ assert .Equal (t , aws .ToString (objects [i ].Key ), outputKeys [i ])
207+ }
208+
209+ // ASSERT all workers and memory returned to pools.
210+ time .Sleep (2 * time .Second )
211+ assert .Equal (t , 100 , len (client .manager .workerPool .channel ))
212+ assert .Equal (t , 100 , len (client .manager .memoryPool .channel ))
213+
214+ // ASSERT that process blocked when all memory secured.
215+ tenByteFile := types.Object {Size : aws .Int64 (10 )}
216+ oneThousandBytesOfFiles := make ([]types.Object , 100 )
217+ for i := 0 ; i < 100 ; i ++ {
218+ oneThousandBytesOfFiles [i ] = tenByteFile
219+ }
220+ client .manager .secureMemory (oneThousandBytesOfFiles )
221+ ch := make (chan chan HydratedFile )
222+ go func () {
223+ ch <- client .GetAllConcurrentlyWithContext (context .Background (), testBucket , "" , objects )
224+ }()
225+
226+ for {
227+ select {
228+ case <- ch :
229+ t .Error ("process was not blocked" )
230+ case <- time .After (time .Second ):
231+ // Timed out as expected
232+ return
233+ }
234+ }
235+ }
236+
126237// go test --run TestS3GetAllConcurrentlyWithContext_Cancel -v
127238func TestS3GetAllConcurrentlyWithContext_Cancel (t * testing.T ) {
128239 // ARRANGE
@@ -175,5 +286,14 @@ func TestS3GetAllConcurrentlyWithContext_Cancel(t *testing.T) {
175286 require .GreaterOrEqual (t , len (collected ), cancelAfter )
176287 // But not all objects should be processed
177288 require .Less (t , len (collected ), len (objects ))
289+ // Pool recovery
290+ require .Eventually (t , func () bool {
291+ return len (client .manager .workerPool .channel ) == 100
292+ }, 5 * time .Second , 10 * time .Millisecond , fmt .Sprintf ("workers pool not recovered, expected %d actual %d" , 100 , len (client .manager .workerPool .channel )))
293+ require .Eventually (t , func () bool {
294+ return len (client .manager .memoryPool .channel ) == 100
295+ }, 5 * time .Second , 10 * time .Millisecond , fmt .Sprintf ("memory pool not recovered, expected %d actual %d" , 100 , len (client .manager .memoryPool .channel )))
296+
178297 })
298+
179299}
0 commit comments