| 1382 | } |
| 1383 | |
| 1384 | func (a *acquireOne) acquireWithCancel(stream proto.DRPCProvisionerDaemon_AcquireJobWithCancelStream) error { |
| 1385 | a.mu.Lock() |
| 1386 | defer a.mu.Unlock() |
| 1387 | a.called++ |
| 1388 | if a.called == 2 { |
| 1389 | close(a.complete) |
| 1390 | } |
| 1391 | if a.called > 1 { |
| 1392 | _, _ = stream.Recv() |
| 1393 | _ = stream.Send(&proto.AcquiredJob{}) |
| 1394 | return nil |
| 1395 | } |
| 1396 | err := stream.Send(a.job) |
| 1397 | // dRPC is racy, and sometimes will return context.Canceled after it has successfully sent the message if we cancel |
| 1398 | // right away, e.g. in unit tests that complete. So, just swallow the error in that case. If we are canceled before |
| 1399 | // the job was acquired, presumably something else in the test will have failed. |
| 1400 | if !xerrors.Is(err, context.Canceled) { |
| 1401 | assert.NoError(a.t, err) |
| 1402 | } |
| 1403 | return nil |
| 1404 | } |
| 1405 | |
| 1406 | func extractInit(t *testing.T) func(s *provisionersdk.Session, r *provisionersdk.InitRequest, canceledOrComplete <-chan struct{}) *sdkproto.InitComplete { |
| 1407 | logger := slogtest.Make(t, nil) |