xref: /aosp_15_r20/external/federated-compute/fcp/tensorflow/host_object.cc (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2019 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include "fcp/tensorflow/host_object.h"
18 
19 #include <utility>
20 
21 #include "fcp/base/monitoring.h"
22 
23 namespace fcp {
24 
25 namespace host_object_internal {
26 
TryLookup(RandomToken token)27 std::optional<std::shared_ptr<void>> HostObjectRegistryImpl::TryLookup(
28     RandomToken token) {
29   std::shared_ptr<void> p = nullptr;
30 
31   {
32     absl::ReaderMutexLock lock{&mutex_};
33     auto it = objects_.find(token);
34     if (it != objects_.end()) {
35       p = it->second;
36     }
37   }
38 
39   if (p == nullptr) {
40     return std::nullopt;
41   } else {
42     return p;
43   }
44 }
45 
Register(RandomToken token,std::shared_ptr<void> p)46 void HostObjectRegistryImpl::Register(RandomToken token,
47                                       std::shared_ptr<void> p) {
48   absl::WriterMutexLock lock{&mutex_};
49   auto r = objects_.insert({token, std::move(p)});
50   FCP_CHECK(r.second)
51       << "An object has already been registered with the provided token";
52 }
53 
Unregister(RandomToken token)54 void HostObjectRegistryImpl::Unregister(RandomToken token) {
55   absl::WriterMutexLock lock{&mutex_};
56   size_t erased = objects_.erase(token);
57   FCP_CHECK(erased == 1)
58       << "An object is not currently registered for the provided token";
59 }
60 
61 }  // namespace host_object_internal
62 
63 }  // namespace fcp
64