|
17 | 17 | #include "Firestore/core/src/firebase/firestore/remote/datastore.h" |
18 | 18 |
|
19 | 19 | #include <unordered_set> |
| 20 | +#include <utility> |
20 | 21 |
|
21 | 22 | #include "Firestore/core/include/firebase/firestore/firestore_errors.h" |
22 | 23 | #include "Firestore/core/src/firebase/firestore/auth/credentials_provider.h" |
23 | 24 | #include "Firestore/core/src/firebase/firestore/auth/token.h" |
24 | 25 | #include "Firestore/core/src/firebase/firestore/core/database_info.h" |
25 | 26 | #include "Firestore/core/src/firebase/firestore/remote/grpc_completion.h" |
| 27 | +#include "Firestore/core/src/firebase/firestore/util/error_apple.h" |
26 | 28 | #include "Firestore/core/src/firebase/firestore/util/executor_libdispatch.h" |
27 | 29 | #include "Firestore/core/src/firebase/firestore/util/hard_assert.h" |
| 30 | +#include "Firestore/core/src/firebase/firestore/util/log.h" |
| 31 | +#include "Firestore/core/src/firebase/firestore/util/statusor.h" |
28 | 32 | #include "absl/memory/memory.h" |
29 | 33 | #include "absl/strings/str_cat.h" |
30 | 34 |
|
|
35 | 39 | using auth::CredentialsProvider; |
36 | 40 | using auth::Token; |
37 | 41 | using core::DatabaseInfo; |
| 42 | +using model::DocumentKey; |
38 | 43 | using util::AsyncQueue; |
39 | 44 | using util::Status; |
| 45 | +using util::StatusOr; |
40 | 46 | using util::internal::Executor; |
41 | 47 | using util::internal::ExecutorLibdispatch; |
42 | 48 |
|
43 | 49 | namespace { |
44 | 50 |
|
| 51 | +const auto kRpcNameLookup = |
| 52 | + "/google.firestore.v1beta1.Firestore/BatchGetDocuments"; |
| 53 | + |
45 | 54 | std::unique_ptr<Executor> CreateExecutor() { |
46 | 55 | auto queue = dispatch_queue_create("com.google.firebase.firestore.rpc", |
47 | 56 | DISPATCH_QUEUE_SERIAL); |
|
56 | 65 | return {grpc_str.begin(), grpc_str.size()}; |
57 | 66 | } |
58 | 67 |
|
| 68 | +void LogGrpcCallFinished(absl::string_view rpc_name, |
| 69 | + GrpcStreamingReader *call, |
| 70 | + const Status &status) { |
| 71 | + LOG_DEBUG("RPC %s completed. Error: %s: %s", rpc_name, status.code(), |
| 72 | + status.error_message()); |
| 73 | + if (bridge::IsLoggingEnabled()) { |
| 74 | + auto headers = |
| 75 | + Datastore::GetWhitelistedHeadersAsString(call->GetResponseHeaders()); |
| 76 | + LOG_DEBUG("RPC %s returned headers (whitelisted): %s", rpc_name, headers); |
| 77 | + } |
| 78 | +} |
| 79 | + |
59 | 80 | } // namespace |
60 | 81 |
|
61 | 82 | Datastore::Datastore(const DatabaseInfo &database_info, |
|
66 | 87 | worker_queue_{worker_queue}, |
67 | 88 | credentials_{credentials}, |
68 | 89 | rpc_executor_{CreateExecutor()}, |
69 | | - serializer_{serializer} { |
| 90 | + serializer_bridge_{serializer} { |
70 | 91 | rpc_executor_->Execute([this] { PollGrpcQueue(); }); |
71 | 92 | } |
72 | 93 |
|
73 | 94 | void Datastore::Shutdown() { |
| 95 | + for (auto &call : lookup_calls_) { |
| 96 | + call->Cancel(); |
| 97 | + } |
| 98 | + lookup_calls_.clear(); |
| 99 | + |
74 | 100 | // `grpc::CompletionQueue::Next` will only return `false` once `Shutdown` has |
75 | 101 | // been called and all submitted tags have been extracted. Without this call, |
76 | 102 | // `rpc_executor_` will never finish. |
|
99 | 125 |
|
100 | 126 | std::shared_ptr<WatchStream> Datastore::CreateWatchStream( |
101 | 127 | id<FSTWatchStreamDelegate> delegate) { |
102 | | - return std::make_shared<WatchStream>(worker_queue_, credentials_, serializer_, |
| 128 | + return std::make_shared<WatchStream>(worker_queue_, credentials_, |
| 129 | + serializer_bridge_.GetSerializer(), |
103 | 130 | &grpc_connection_, delegate); |
104 | 131 | } |
105 | 132 |
|
106 | 133 | std::shared_ptr<WriteStream> Datastore::CreateWriteStream( |
107 | 134 | id<FSTWriteStreamDelegate> delegate) { |
108 | | - return std::make_shared<WriteStream>(worker_queue_, credentials_, serializer_, |
| 135 | + return std::make_shared<WriteStream>(worker_queue_, credentials_, |
| 136 | + serializer_bridge_.GetSerializer(), |
109 | 137 | &grpc_connection_, delegate); |
110 | 138 | } |
111 | 139 |
|
| 140 | +void Datastore::LookupDocuments( |
| 141 | + const std::vector<DocumentKey> &keys, |
| 142 | + FSTVoidMaybeDocumentArrayErrorBlock completion) { |
| 143 | + grpc::ByteBuffer message = serializer_bridge_.ToByteBuffer( |
| 144 | + serializer_bridge_.CreateLookupRequest(keys)); |
| 145 | + |
| 146 | + ResumeRpcWithCredentials( |
| 147 | + [this, message, completion](const StatusOr<Token> &maybe_credentials) { |
| 148 | + if (!maybe_credentials.ok()) { |
| 149 | + completion(nil, util::MakeNSError(maybe_credentials.status())); |
| 150 | + return; |
| 151 | + } |
| 152 | + LookupDocumentsWithCredentials(maybe_credentials.ValueOrDie(), message, |
| 153 | + completion); |
| 154 | + }); |
| 155 | +} |
| 156 | + |
| 157 | +void Datastore::LookupDocumentsWithCredentials( |
| 158 | + const Token &token, |
| 159 | + const grpc::ByteBuffer &message, |
| 160 | + FSTVoidMaybeDocumentArrayErrorBlock completion) { |
| 161 | + lookup_calls_.push_back(grpc_connection_.CreateStreamingReader( |
| 162 | + kRpcNameLookup, token, std::move(message))); |
| 163 | + GrpcStreamingReader *call = lookup_calls_.back().get(); |
| 164 | + |
| 165 | + call->Start([this, call, completion]( |
| 166 | + const StatusOr<std::vector<grpc::ByteBuffer>> &result) { |
| 167 | + OnLookupDocumentsResponse(call, result, completion); |
| 168 | + RemoveGrpcCall(call); |
| 169 | + }); |
| 170 | +} |
| 171 | + |
| 172 | +void Datastore::OnLookupDocumentsResponse( |
| 173 | + GrpcStreamingReader *call, |
| 174 | + const StatusOr<std::vector<grpc::ByteBuffer>> &result, |
| 175 | + FSTVoidMaybeDocumentArrayErrorBlock completion) { |
| 176 | + LogGrpcCallFinished("BatchGetDocuments", call, result.status()); |
| 177 | + HandleCallStatus(result.status()); |
| 178 | + |
| 179 | + if (!result.ok()) { |
| 180 | + completion(nil, util::MakeNSError(result.status())); |
| 181 | + return; |
| 182 | + } |
| 183 | + |
| 184 | + Status parse_status; |
| 185 | + std::vector<grpc::ByteBuffer> responses = std::move(result).ValueOrDie(); |
| 186 | + NSArray<FSTMaybeDocument *> *docs = |
| 187 | + serializer_bridge_.MergeLookupResponses(responses, &parse_status); |
| 188 | + if (parse_status.ok()) { |
| 189 | + completion(docs, nil); |
| 190 | + } else { |
| 191 | + completion(nil, util::MakeNSError(parse_status)); |
| 192 | + } |
| 193 | +} |
| 194 | + |
| 195 | +void Datastore::ResumeRpcWithCredentials(const OnCredentials &on_credentials) { |
| 196 | + // Auth may outlive Firestore |
| 197 | + std::weak_ptr<Datastore> weak_this{shared_from_this()}; |
| 198 | + |
| 199 | + credentials_->GetToken( |
| 200 | + [weak_this, on_credentials](const StatusOr<Token> &result) { |
| 201 | + auto strong_this = weak_this.lock(); |
| 202 | + if (!strong_this) { |
| 203 | + return; |
| 204 | + } |
| 205 | + |
| 206 | + strong_this->worker_queue_->EnqueueRelaxed( |
| 207 | + [weak_this, result, on_credentials] { |
| 208 | + auto strong_this = weak_this.lock(); |
| 209 | + if (!strong_this) { |
| 210 | + return; |
| 211 | + } |
| 212 | + |
| 213 | + on_credentials(result); |
| 214 | + }); |
| 215 | + }); |
| 216 | +} |
| 217 | + |
| 218 | +void Datastore::HandleCallStatus(const Status &status) { |
| 219 | + if (status.code() == FirestoreErrorCode::Unauthenticated) { |
| 220 | + credentials_->InvalidateToken(); |
| 221 | + } |
| 222 | +} |
| 223 | + |
| 224 | +void Datastore::RemoveGrpcCall(GrpcStreamingReader *to_remove) { |
| 225 | + auto found = std::find_if( |
| 226 | + lookup_calls_.begin(), lookup_calls_.end(), |
| 227 | + [to_remove](const std::unique_ptr<GrpcStreamingReader> &call) { |
| 228 | + return call.get() == to_remove; |
| 229 | + }); |
| 230 | + HARD_ASSERT(found != lookup_calls_.end(), "Missing gRPC call"); |
| 231 | + lookup_calls_.erase(found); |
| 232 | +} |
| 233 | + |
112 | 234 | std::string Datastore::GetWhitelistedHeadersAsString( |
113 | 235 | const GrpcStream::MetadataT &headers) { |
114 | 236 | static std::unordered_set<std::string> whitelist = { |
|
0 commit comments