| 24 | AllocatorRetry::AllocatorRetry() : env_(Env::Default()) {} |
| 25 | |
| 26 | void* AllocatorRetry::AllocateRaw( |
| 27 | std::function<void*(size_t alignment, size_t num_bytes, |
| 28 | bool verbose_failure)> |
| 29 | alloc_func, |
| 30 | int max_millis_to_wait, size_t alignment, size_t num_bytes) { |
| 31 | if (num_bytes == 0) { |
| 32 | return nullptr; |
| 33 | } |
| 34 | uint64 deadline_micros = 0; |
| 35 | bool first = true; |
| 36 | void* ptr = nullptr; |
| 37 | while (ptr == nullptr) { |
| 38 | ptr = alloc_func(alignment, num_bytes, false); |
| 39 | if (ptr == nullptr) { |
| 40 | uint64 now = env_->NowMicros(); |
| 41 | if (first) { |
| 42 | deadline_micros = now + max_millis_to_wait * 1000; |
| 43 | first = false; |
| 44 | } |
| 45 | if (now < deadline_micros) { |
| 46 | mutex_lock l(mu_); |
| 47 | WaitForMilliseconds(&l, &memory_returned_, |
| 48 | (deadline_micros - now) / 1000); |
| 49 | } else { |
| 50 | return alloc_func(alignment, num_bytes, true); |
| 51 | } |
| 52 | } |
| 53 | } |
| 54 | return ptr; |
| 55 | } |
| 56 | |
| 57 | } // namespace tensorflow |
nothing calls this directly
no test coverage detected