const ServableId& id, const Status load_status = Status::OK()) {

std::unique_ptr<Loader> loader(new FakeLoader(id.version, load_status));

return CreateServableData(id, std::move(loader));

protected:

BasicManagerTest()

: thread_pool_sizes_(GetParam()),

servable_event_bus_(EventBus<ServableState>::CreateEventBus()),

servable_state_monitor_(servable_event_bus_.get()) {

BasicManager::Options options;

options.num_load_threads = thread_pool_sizes_.num_load_threads;

options.num_unload_threads = thread_pool_sizes_.num_unload_threads;

options.servable_event_bus = servable_event_bus_.get();

options.max_num_load_retries = 10;

options.load_retry_interval_micros = 0;

TF_CHECK_OK(BasicManager::Create(std::move(options), &basic_manager_));

}

void SetUp() override {

// We load the manager with two different servable streams, each with two

// versions 0 and 1.

std::set<ServableId> loaded_servables;

for (const char* servable_name : {kServableName, kServableName2}) {

for (int i = 1; i <= kNumVersionsPerServable; ++i) {

const ServableId id = {servable_name, i};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

loaded_servables.insert(id);

}

}

for (const ServableId& loaded_servable : loaded_servables) {

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, loaded_servable,

{ServableState::ManagerState::kAvailable});

}

}

ThreadPoolSizes thread_pool_sizes_;

std::shared_ptr<EventBus<ServableState>> servable_event_bus_;

ServableStateMonitor servable_state_monitor_;

std::unique_ptr<BasicManager> basic_manager_;

ServableHandle<int64> handle;

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Latest(strings::StrCat(kServableName, "missing")),

&handle);

ASSERT_FALSE(status.ok()) << status;

EXPECT_EQ(error::NOT_FOUND, status.code());

// This version is missing.

const int64 missing_version = 100;

ServableHandle<int64> handle;

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Specific(kServableName, missing_version), &handle);

ASSERT_FALSE(status.ok()) << status;

EXPECT_EQ(error::NOT_FOUND, status.code());

const ServableId id = {kServableName, kNumVersionsPerServable + 1};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

ServableHandle<int64> handle;

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Latest(kServableName), &handle);

TF_ASSERT_OK(status);

EXPECT_EQ(kNumVersionsPerServable + 1, *handle);

const ServableId id = {"banana", 42};

TF_ASSERT_OK(basic_manager_->ManageServable(CreateServable(id)));

EXPECT_FALSE(basic_manager_->ManageServable(CreateServable(id)).ok());

const ServableId id = {kServableName3, 1};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

ServableHandle<int64> handle;

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Latest(kServableName3), &handle);

TF_ASSERT_OK(status);

EXPECT_EQ(1, *handle);

EXPECT_EQ(id, handle.id());

const ServableId id = {kServableName3, 1};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_EXPECT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

EXPECT_FALSE(basic_manager_->StopManagingServable(id).ok());

const ServableId id = {kServableName3, 1};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_EXPECT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

basic_manager_->UnloadServable(

id, [](const Status& status) { TF_EXPECT_OK(status); });

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

const optional<ServableStateSnapshot<>> snapshot =

basic_manager_->GetManagedServableStateSnapshot(id);

EXPECT_EQ(LoaderHarness::State::kDisabled, snapshot->state);

const ServableState expected_state = {id, ServableState::ManagerState::kEnd,

Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(expected_state));

TF_ASSERT_OK(basic_manager_->StopManagingServable(id));

EXPECT_FALSE(basic_manager_->GetManagedServableStateSnapshot(id));

const ServableId id = {kServableName, 7};

const Status error_status = errors::Internal("An error.");

std::unique_ptr<Loader> loader(new FakeLoader(7, error_status));

basic_manager_->ManageServable({id, std::move(loader)});

basic_manager_->LoadServable(id, [error_status](const Status& status) {

EXPECT_EQ(error_status, status);

});

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

const optional<ServableStateSnapshot<>> snapshot =

basic_manager_->GetManagedServableStateSnapshot(id);

EXPECT_EQ(LoaderHarness::State::kError, snapshot->state);

const ServableState expected_error_state = {

id, ServableState::ManagerState::kEnd, error_status};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(expected_error_state));

ServableHandle<int64> handle;

const ServableId id = {kServableName2, 1};

const Status status =

basic_manager_->GetServableHandle(ServableRequest::FromId(id), &handle);

TF_ASSERT_OK(status);

EXPECT_EQ(1, *handle);

EXPECT_EQ(id, handle.id());

// Using kServableName3 which doesn't have any servables loaded in the

// manager, as opposed to kServableName which already has 2 loaded.

const ServableId id0 = {kServableName3, 0};

// Servable is int64 with value 0.

basic_manager_->ManageServable(CreateServable(id0));

basic_manager_->LoadServable(

id0, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id0, {ServableState::ManagerState::kAvailable});

test_util::MockLoader* notify_to_unload = new NiceMock<test_util::MockLoader>;

// Don't make it const otherwise servable types will mismatch: const int64 vs

// int64.

int64 servable = 1;

ON_CALL(*notify_to_unload, servable())

.WillByDefault(Return(AnyPtr(&servable)));

ON_CALL(*notify_to_unload, EstimateResources(_))

.WillByDefault(Return(Status::OK()));

ON_CALL(*notify_to_unload, Load()).WillByDefault(Return(Status::OK()));

const ServableId id1 = {kServableName3, 1};

basic_manager_->ManageServable(

{id1, std::unique_ptr<Loader>(notify_to_unload)});

basic_manager_->LoadServable(

id1, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id1, {ServableState::ManagerState::kAvailable});

// We have loaded both versions 0 and 1 of kServableName3, so the latest

// handle should be that of v1.

{

ServableHandle<int64> handle;

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Latest(kServableName3), &handle);

TF_ASSERT_OK(status);

EXPECT_EQ(id1, handle.id());

}

// We will now try to unload v1, but we only allow it to move out from kReady

// state, and not complete the unload. Also, after it moves out from kReady,

// the serving map is also updated, so v0 would be the latest.

Notification unload_started;

Notification finish_unload;

EXPECT_CALL(*notify_to_unload, Unload()).WillOnce(Invoke([&]() {

unload_started.Notify();

finish_unload.WaitForNotification();

}));

Notification unload_finished;

std::unique_ptr<Thread> unload_last_servable(

Env::Default()->StartThread({}, "UnloadLastServable", [&]() {

basic_manager_->UnloadServable(id1, [&](const Status& status) {

TF_EXPECT_OK(status);

unload_finished.Notify();

});

}));

unload_started.WaitForNotification();

// Servable map should just have {kServableName3, 0} at this point.

{

ServableHandle<int64> handle;

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Latest(kServableName3), &handle);

TF_EXPECT_OK(status);

EXPECT_EQ(id0, handle.id());

}

finish_unload.Notify();

// We have to ensure that the unload has finished completely, otherwise the

// address of the notifications could be invalid in the load when we exit from

// this scope.

unload_finished.WaitForNotification();

const std::vector<ServableId> expected_before = {{kServableName, 1},

{kServableName, 2},

{kServableName2, 1},

{kServableName2, 2}};

EXPECT_THAT(basic_manager_->ListAvailableServableIds(),

UnorderedElementsAreArray(expected_before));

// Set stream kServableName to have servables 7 and unload 0 & 1, but 7 errors

// on load, so never moves to a loaded state.

const ServableId id = {kServableName, 7};

std::unique_ptr<Loader> loader(

new FakeLoader(7, errors::Internal("An error.")));

basic_manager_->ManageServable(CreateServableData(id, std::move(loader)));

basic_manager_->LoadServable(id, [](const Status& status) {

EXPECT_EQ(errors::Internal("An error."), status);

});

basic_manager_->UnloadServable(

{kServableName, 1}, [](const Status& status) { TF_ASSERT_OK(status); });

basic_manager_->UnloadServable(

{kServableName, 2}, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_,

{kServableName, 1},

{ServableState::ManagerState::kEnd});

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_,

{kServableName, 2},

{ServableState::ManagerState::kEnd});

const std::vector<ServableId> expected_after = {{kServableName2, 1},

{kServableName2, 2}};

EXPECT_THAT(basic_manager_->ListAvailableServableIds(),

UnorderedElementsAreArray(expected_after));

// Scoped to destruct handles at the end of it.

{

const std::map<ServableId, ServableHandle<int64>> handles_before =

basic_manager_->GetAvailableServableHandles<int64>();

ASSERT_EQ(kNumVersionsPerServable * 2, handles_before.size());

const std::vector<ServableId> expected_ids_before = {{kServableName, 1},

{kServableName, 2},

{kServableName2, 1},

{kServableName2, 2}};

for (const ServableId& expected_id : expected_ids_before) {

const auto found_it = handles_before.find(expected_id);

ASSERT_TRUE(found_it != handles_before.end());

EXPECT_EQ(expected_id.version, *found_it->second);

}

}

// Set stream kServableName to have servables 7 and unload 0 & 1, but 7 errors

// on load, so never moves to a loaded state.

const ServableId id = {kServableName, 7};

std::unique_ptr<Loader> loader(

new FakeLoader(7, errors::Internal("An error.")));

basic_manager_->ManageServable(CreateServableData(id, std::move(loader)));

basic_manager_->LoadServable(id, [](const Status& status) {

EXPECT_EQ(errors::Internal("An error."), status);

});

basic_manager_->UnloadServable(

{kServableName, 1}, [](const Status& status) { TF_ASSERT_OK(status); });

basic_manager_->UnloadServable(

{kServableName, 2}, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_,

{kServableName, 1},

{ServableState::ManagerState::kEnd});

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_,

{kServableName, 2},

{ServableState::ManagerState::kEnd});

{

const std::map<ServableId, ServableHandle<int64>> handles_after =

basic_manager_->GetAvailableServableHandles<int64>();

ASSERT_EQ(kNumVersionsPerServable, handles_after.size());

const std::vector<ServableId> expected_ids_after = {{kServableName2, 1},

{kServableName2, 2}};

for (const ServableId& expected_id : expected_ids_after) {

const auto found_it = handles_after.find(expected_id);

ASSERT_TRUE(found_it != handles_after.end());

EXPECT_EQ(expected_id.version, *found_it->second);

}

}

const std::map<ServableId, ServableHandle<int>> wrong_type_handles =

basic_manager_->GetAvailableServableHandles<int>();

EXPECT_EQ(0, wrong_type_handles.size());

GetManagedServableStateSnapshotWithoutAdditionalState) {

const std::vector<ServableStateSnapshot<>> expected = {

{{kServableName, 1}, LoaderHarness::State::kReady, {}},

{{kServableName, 2}, LoaderHarness::State::kReady, {}}};

EXPECT_THAT(basic_manager_->GetManagedServableStateSnapshots(kServableName),

UnorderedElementsAreArray(expected));

// Check servable state snapshot corresponding to a servable-id that is in

// ready state.

const ServableId id_ready = {kServableName, 1};

const optional<ServableStateSnapshot<>> actual_ready_snapshot =

basic_manager_->GetManagedServableStateSnapshot(id_ready);

EXPECT_TRUE(actual_ready_snapshot);

const ServableStateSnapshot<> expected_ready_snapshot = {

id_ready, LoaderHarness::State::kReady, {}};

EXPECT_EQ(actual_ready_snapshot, expected_ready_snapshot);

// Check servable state snapshot corresponding to a servable-id that is not

// managed by the basic-manager.

const ServableId id_notmanaged = {kServableName, 8};

EXPECT_FALSE(basic_manager_->GetManagedServableStateSnapshot(id_notmanaged));

basic_manager_->ManageServableWithAdditionalState(

CreateServable({kServableName3, 0}), std::unique_ptr<int>(new int(0)));

basic_manager_->ManageServableWithAdditionalState(

CreateServable({kServableName3, 1}), std::unique_ptr<int>(new int(1)));

const std::vector<ServableStateSnapshot<int>> expected = {

{{kServableName3, 0}, LoaderHarness::State::kNew, {0}},

{{kServableName3, 1}, LoaderHarness::State::kNew, {1}}};

EXPECT_THAT(

basic_manager_->GetManagedServableStateSnapshots<int>(kServableName3),

UnorderedElementsAreArray(expected));

const ServableId id = {kServableName, 1};

std::unique_ptr<Loader> first_loader(

new FakeLoader(1, errors::Internal("An error.")));

basic_manager_->ManageServable(

CreateServableData(id, std::move(first_loader)));

// Different servable returned.

std::unique_ptr<Loader> second_loader(

new FakeLoader(2, errors::Internal("An error.")));

basic_manager_->ManageServable(

CreateServableData(id, std::move(second_loader)));

ServableHandle<int64> handle;

TF_ASSERT_OK(basic_manager_->GetServableHandle(

ServableRequest::Specific(kServableName, 1), &handle));

EXPECT_EQ(1, *handle);

const ServableId id = {kServableName, 3};

basic_manager_->ManageServable(

ServableData<std::unique_ptr<Loader>>(id, errors::Unknown("error")));

ServableHandle<int64> handle;

Status status = basic_manager_->GetServableHandle(

ServableRequest::Specific(kServableName, 3), &handle);

EXPECT_FALSE(status.ok()) << status;

basic_manager_->LoadServable(

id, [](const Status& status) { EXPECT_FALSE(status.ok()) << status; });

status = basic_manager_->GetServableHandle(

ServableRequest::Specific(kServableName, 3), &handle);

EXPECT_FALSE(status.ok()) << status;

const ServableId id = {kServableName, 7};

basic_manager_->ManageServable(

CreateServableData(id, std::unique_ptr<Loader>(new FakeLoader(7))));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

std::unique_ptr<ServableHandle<int64>> latest_handle(

new ServableHandle<int64>());

const Status status = basic_manager_->GetServableHandle(

ServableRequest::Latest(kServableName), latest_handle.get());

TF_ASSERT_OK(status);

EXPECT_EQ(7, **latest_handle);

Notification done_unload_servable;

std::unique_ptr<Thread> unload_servable(

Env::Default()->StartThread({}, "UnloadServable", [&]() {

// Unload the servable.

basic_manager_->UnloadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kEnd});

basic_manager_->StopManagingServable(id);

// The servable has been deleted in this thread if there is no

// thread-pool for load/unload.

if (thread_pool_sizes_.num_load_threads == 0) {

EXPECT_TRUE(FakeLoader::was_deleted_in_this_thread());

}

done_unload_servable.Notify();

}));

// This will unblock the UnloadServable.

latest_handle.reset();

done_unload_servable.WaitForNotification();

// The servable wasn't deleted in this thread.

ASSERT_FALSE(FakeLoader::was_deleted_in_this_thread());

const ServableId id = {kServableName, 3};

std::unique_ptr<int> state(new int(1));

basic_manager_->ManageServableWithAdditionalState(CreateServable(id),

std::move(state));

EXPECT_EQ(1, *basic_manager_->GetAdditionalServableState<int>(id));

EXPECT_EQ(nullptr, basic_manager_->GetAdditionalServableState<float>(id));

const ServableId id = {kServableName, 3};

basic_manager_->ManageServable(CreateServable(id));

// Will return nullptr when there is no metadata set.

EXPECT_EQ(nullptr, basic_manager_->GetAdditionalServableState<int>(id));

EXPECT_EQ(nullptr, basic_manager_->GetAdditionalServableState<float>(id));

const ServableId id = {kServableName, 3};

basic_manager_->LoadServable(id, [](const Status& status) {

EXPECT_FALSE(status.ok());

EXPECT_EQ(error::NOT_FOUND, status.code());

EXPECT_THAT(status.error_message(), HasSubstr("is not being managed"));

});

const ServableId id = {kServableName, 3};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

basic_manager_->LoadServable(id, [](const Status& status) {

EXPECT_FALSE(status.ok());

EXPECT_EQ(error::FAILED_PRECONDITION, status.code());

EXPECT_THAT(status.error_message(), HasSubstr("Duplicate load request"));

});

const ServableId id = {kServableName, 3};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

basic_manager_->UnloadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

basic_manager_->UnloadServable(id, [](const Status& status) {

EXPECT_FALSE(status.ok());

EXPECT_EQ(error::FAILED_PRECONDITION, status.code());

EXPECT_THAT(status.error_message(),

HasSubstr("unload already requested/ongoing"));

});

const ServableId id = {kServableName, 3};

basic_manager_->UnloadServable(id, [](const Status& status) {

EXPECT_FALSE(status.ok());

EXPECT_EQ(error::NOT_FOUND, status.code());

EXPECT_THAT(status.error_message(), HasSubstr("is not being managed"));

});

const ServableId id = {kServableName, 3};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->UnloadServable(id, [](const Status& status) {

EXPECT_FALSE(status.ok());

EXPECT_EQ(error::FAILED_PRECONDITION, status.code());

EXPECT_THAT(status.error_message(), HasSubstr("Servable not loaded"));

});

const ServableId id = {kServableName, 3};

basic_manager_->ManageServable(

ServableData<std::unique_ptr<Loader>>(id, errors::Unknown("error")));

const ServableState expected_published_state = {

id, ServableState::ManagerState::kEnd, errors::Unknown("error")};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(expected_published_state));

const ServableId id = {kServableName, 7};

std::unique_ptr<Loader> loader(

new FakeLoader(7, errors::Internal("Error on load.")));

basic_manager_->ManageServable({id, std::move(loader)});

const ServableState start_state = {id, ServableState::ManagerState::kStart,

Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(start_state));

basic_manager_->LoadServable(id, [](const Status& status) {});

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

const ServableState error_state = {id, ServableState::ManagerState::kEnd,

errors::Internal("Error on load.")};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(error_state));

const ServableId id = {kServableName, 7};

test_util::MockLoader* loader = new NiceMock<test_util::MockLoader>();

basic_manager_->ManageServable({id, std::unique_ptr<Loader>(loader)});

const ServableState start_state = {id, ServableState::ManagerState::kStart,

Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(start_state));

Notification load_called;

Notification load_continue;

EXPECT_CALL(*loader, Load()).WillOnce(InvokeWithoutArgs([&]() {

load_called.Notify();

load_continue.WaitForNotification();

return Status::OK();

}));

std::unique_ptr<Thread> load_thread(

Env::Default()->StartThread(ThreadOptions(), "LoadThread", [&]() {

basic_manager_->LoadServable(id, [](const Status& status) {});

}));

load_called.WaitForNotification();

const ServableState loading_state = {

id, ServableState::ManagerState::kLoading, Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(loading_state));

load_continue.Notify();

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

const ServableState available_state = {

id, ServableState::ManagerState::kAvailable, Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(available_state));

Notification unload_called;

Notification unload_continue;

EXPECT_CALL(*loader, Unload())

.WillOnce(Invoke([&]() {

unload_called.Notify();

unload_continue.WaitForNotification();

}));

// Scoped to ensure UnloadServable() is scheduled.

std::unique_ptr<Thread> unload_thread(

Env::Default()->StartThread(ThreadOptions(), "UnloadThread", [&]() {

basic_manager_->UnloadServable(id, [](const Status& status) {});

}));

unload_called.WaitForNotification();

const ServableState unloading_state = {

id, ServableState::ManagerState::kUnloading, Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(unloading_state));

unload_continue.Notify();

WaitUntilServableManagerStateIsOneOf(servable_state_monitor_, id,

{ServableState::ManagerState::kEnd});

const ServableState end_state = {id, ServableState::ManagerState::kEnd,

Status::OK()};

EXPECT_THAT(*servable_state_monitor_.GetState(id),

EqualsServableState(end_state));

BasicManager::Options options;

// Single threaded execution.

options.num_load_threads = 0;

// No event bus.

options.servable_event_bus = nullptr;

std::unique_ptr<BasicManager> manager;

TF_ASSERT_OK(BasicManager::Create(std::move(options), &manager));

const ServableId id = {kServableName, 7};

std::unique_ptr<Loader> loader(new FakeLoader(7));

manager->ManageServable({id, std::move(loader)});

manager->LoadServable(id, [](const Status& status) { TF_ASSERT_OK(status); });

manager->UnloadServable(id,

[](const Status& status) { TF_ASSERT_OK(status); });

std::set<ServableId> servables;

// Scoped so that all loads can be scheduled before proceeding.

{

ThreadPoolExecutor load_executor(Env::Default(), "LoadServables",

kNumThreads);

for (int i = 0; i < 20; ++i) {

const ServableId id = {kServableName3, i};

servables.insert(id);

load_executor.Schedule([this, id, &servables]() {

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

});

}

}

// At this point, all loads may not have completed, so we wait for them.

for (const ServableId& servable : servables) {

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, servable,

{ServableState::ManagerState::kAvailable});

}

{

ThreadPoolExecutor unload_executor(Env::Default(), "UnloadServables",

kNumThreads);

// Doing in reverse.

for (int i = 19; i >= 0; --i) {

unload_executor.Schedule([this, i]() {

const ServableId id = {kServableName3, i};

basic_manager_->UnloadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

});

}

}

ThreadPoolExecutor executor(Env::Default(), "InterleavedLoadsAndUnloads",

kNumThreads);

for (int i = 0; i < 20; ++i) {

executor.Schedule([this, i]() {

const ServableId id = {kServableName3, i};

basic_manager_->ManageServable(CreateServable(id));

Notification load_done;

basic_manager_->LoadServable(id, [&load_done](const Status& status) {

TF_ASSERT_OK(status);

load_done.Notify();

});

load_done.WaitForNotification();

basic_manager_->UnloadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

});

}

protected:

SetNumLoadThreadsBasicManagerTest() {

BasicManager::Options options;

options.num_load_threads = 0;

options.max_num_load_retries = 10;

options.load_retry_interval_micros = 0;

TF_CHECK_OK(BasicManager::Create(std::move(options), &basic_manager_));

}

std::unique_ptr<BasicManager> basic_manager_;

test_util::BasicManagerTestAccess manager_test_access(basic_manager_.get());

manager_test_access.SetNumLoadThreads(2);

EXPECT_EQ(2, manager_test_access.num_load_threads());

const auto load_done_fn = [&](const Status& status) {

TF_ASSERT_OK(status);

// Tests whether the threadpools are actually swapped in

// SetNumLoadThreads().

static thread_local int per_thread_load_ctr = 0;

++per_thread_load_ctr;

EXPECT_EQ(1, per_thread_load_ctr);

};

const ServableId id0 = {kServableName3, 0};

basic_manager_->ManageServable(CreateServable(id0));

basic_manager_->LoadServable(id0, load_done_fn);

manager_test_access.SetNumLoadThreads(0);

EXPECT_EQ(0, manager_test_access.num_load_threads());

const ServableId id1 = {kServableName3, 1};

basic_manager_->ManageServable(CreateServable(id1));

basic_manager_->LoadServable(id1, load_done_fn);

// Force the manager to finish before deleting the notifications.

basic_manager_.reset();

test_util::BasicManagerTestAccess manager_test_access(basic_manager_.get());

manager_test_access.SetNumLoadThreads(1);

EXPECT_EQ(1, manager_test_access.num_load_threads());

std::set<string> data_race_set;

const auto data_race_fn = [&](const Status& status) {

// This line will cause a data race if both the loads happen simultaneously

// on different threads. This will be caught by the ThreadSanitizer, causing

// the test to fail.

data_race_set.insert("string");

};

const ServableId id0 = {kServableName3, 0};

basic_manager_->ManageServable(CreateServable(id0));

Notification notify_for_setting;

Notification continue_load;

basic_manager_->LoadServable(id0, [&](const Status& status) {

notify_for_setting.Notify();

continue_load.WaitForNotification();

data_race_fn(status);

});

{

ThreadPoolExecutor executor(Env::Default(), "SetNumLoadThreads",

kNumThreads);

executor.Schedule([&]() {

notify_for_setting.WaitForNotification();

manager_test_access.SetNumLoadThreads(1);

EXPECT_EQ(1, manager_test_access.num_load_threads());

});

executor.Schedule([&]() {

const ServableId id1 = {kServableName3, 1};

basic_manager_->ManageServable(CreateServable(id1));

continue_load.Notify();

basic_manager_->LoadServable(

id1, [&](const Status& status) { data_race_fn(status); });

});

}

// Force the manager to finish before deleting the notifications.

basic_manager_.reset();

test_util::BasicManagerTestAccess manager_test_access(basic_manager_.get());

const uint32 prev_num_load_threads = manager_test_access.num_load_threads();

manager_test_access.SetNumLoadThreads(32);

EXPECT_EQ(32, manager_test_access.num_load_threads());

{

ThreadPoolExecutor executor(Env::Default(), "FirstThreadPoolLoads",

kNumThreads);

for (int i = 0; i < 20; ++i) {

executor.Schedule([this, i]() {

const ServableId id = {kServableName3, i};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

// We don't wait for load to be done here because we want to test that

// SetNumLoadThreads() waits properly till all queued loads are

// finished. If a queued load hasn't been finished the corresponding

// UnloadServable() will fail.

});

}

}

manager_test_access.SetNumLoadThreads(prev_num_load_threads);

EXPECT_EQ(prev_num_load_threads, manager_test_access.num_load_threads());

{

ThreadPoolExecutor executor(Env::Default(), "Unloads", kNumThreads);

for (int i = 0; i < 20; ++i) {

executor.Schedule([this, i]() {

const ServableId id = {kServableName3, i};

basic_manager_->UnloadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

});

}

}

const ServableId id = {kServableName3, 0};

mutex status_mu;

std::vector<Status> statuses(4);

{

ThreadPoolExecutor load_executor(Env::Default(), "LoadServables",

kNumThreads);

for (int i = 0; i < 4; ++i) {

load_executor.Schedule([this, id, i, &statuses, &status_mu]() {

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [i, &statuses, &status_mu](const Status& status) {

mutex_lock l(status_mu);

statuses[i] = status;

});

});

}

}

// At this point, all loads may not have completed. Deleting BasicManager

// would wait for all the scheduled loads to complete before deleting it.

basic_manager_.reset();

int num_status_ok = 0;

for (int i = 0; i < 4; ++i) {

mutex_lock l(status_mu);

if (!statuses[i].ok()) {

EXPECT_EQ(error::FAILED_PRECONDITION, statuses[i].code());

EXPECT_THAT(statuses[i].error_message(),

HasSubstr("Duplicate load request"));

} else {

++num_status_ok;

}

}

EXPECT_EQ(1, num_status_ok);

const ServableId id = {kServableName3, 0};

basic_manager_->ManageServable(CreateServable(id));

basic_manager_->LoadServable(

id, [](const Status& status) { TF_ASSERT_OK(status); });

// At this point, all loads may not have completed, so we wait for them.

WaitUntilServableManagerStateIsOneOf(

servable_state_monitor_, id, {ServableState::ManagerState::kAvailable});

mutex status_mu;

std::vector<Status> statuses(4);

{

ThreadPoolExecutor load_executor(Env::Default(), "LoadServables",

kNumThreads);

for (int i = 0; i < 4; ++i) {

load_executor.Schedule([this, id, i, &statuses, &status_mu]() {

basic_manager_->UnloadServable(

id, [i, &statuses, &status_mu](const Status& status) {

mutex_lock l(status_mu);

statuses[i] = status;

});

});

}

}

// At this point, all unloads may not have completed. Deleting BasicManager

// would wait for all the scheduled unloads to complete before deleting it.

basic_manager_.reset();

int num_status_ok = 0;