Skip to content

Commit e085cb3

Browse files
wenbozhutensorflow-copybara
authored andcommitted
Refined the httpserver flow-control API for streaming writes.
PiperOrigin-RevId: 294679764
1 parent ee66538 commit e085cb3

File tree

3 files changed

+43
-18
lines changed

3 files changed

+43
-18
lines changed

tensorflow_serving/util/net_http/server/internal/evhttp_request.cc

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -346,9 +346,10 @@ void EvHTTPRequest::PartialReply() {
346346
NET_LOG(FATAL, "PartialReplyWithStatus not implemented.");
347347
}
348348

349-
void EvHTTPRequest::PartialReplyWithFlushCallback(
350-
std::function<void(bool result)> callback) {
349+
ServerRequestInterface::CallbackStatus
350+
EvHTTPRequest::PartialReplyWithFlushCallback(std::function<void()> callback) {
351351
NET_LOG(FATAL, "PartialReplyWithStatus not implemented.");
352+
return CallbackStatus::NOT_SCHEDULED;
352353
}
353354

354355
void EvHTTPRequest::ReplyWithStatus(HTTPStatusCode status) {

tensorflow_serving/util/net_http/server/internal/evhttp_request.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,8 @@ class EvHTTPRequest final : public ServerRequestInterface {
9696
void PartialReplyWithStatus(HTTPStatusCode status) override;
9797
void PartialReply() override;
9898

99-
void PartialReplyWithFlushCallback(
100-
std::function<void(bool result)> callback) override;
99+
CallbackStatus PartialReplyWithFlushCallback(
100+
std::function<void()> callback) override;
101101

102102
void ReplyWithStatus(HTTPStatusCode status) override;
103103
void Reply() override;

tensorflow_serving/util/net_http/server/public/server_request_interface.h

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,23 @@ class ServerRequestInterface {
118118
virtual void AppendResponseHeader(absl::string_view header,
119119
absl::string_view value) = 0;
120120

121+
// The IO status of a request or response body.
122+
enum class BodyStatus {
123+
// The body hasn't been completely read or written.
124+
PENDING = 0,
125+
// The body has been completely read or written or when there is no body.
126+
COMPLETE = 1,
127+
// The transport has reported a failure and the request should be aborted.
128+
FAILED = 2,
129+
};
130+
131+
// This serves as the return value type for callbacks that may be
132+
// skipped for optimization reasons
133+
enum class CallbackStatus {
134+
NOT_SCHEDULED = 0,
135+
SCHEDULED = 1,
136+
};
137+
121138
// Sends headers and/or any buffered response body data to the client.
122139
// Assumes 200 if status is not specified.
123140
// If called for the first time, all the response headers will be sent
@@ -130,34 +147,41 @@ class ServerRequestInterface {
130147
virtual void PartialReply() = 0;
131148

132149
// Similar to PartialReply() but with an on_flush callback which will be
133-
// invoked when the response data has been completely flushed by the transport
134-
// or when the write fails due to transport errors. This allows the handler to
135-
// respect transport-provided flow-control in writing data to the peer.
150+
// invoked when the response data has been completely flushed by the
151+
// transport. This allows the handler to apply transport-provided flow-control
152+
// in writing data to the peer.
153+
//
154+
// Returns SCHEDULED if the callback will be invoked asynchronously after
155+
// this method returns. Until the callback is invoked, the request object
156+
// should not be accessed by the handler.
136157
//
137-
// Until the callback is invoked, the request object should not be accessed
138-
// by the handler after this method is called. If a failure occurs,
139-
// the callback will be passed a false value, and the request should be
140-
// aborted thereafter by the handler.
158+
// Returns NOT_SCHEDULED if data is already flushed when this method returns
159+
// or when the request should be aborted due to transport failures.
141160
//
142-
// The callback may be invoked immediately but never from the current
143-
// call stack.
144-
virtual void PartialReplyWithFlushCallback(
145-
std::function<void(bool result)> callback) = 0;
161+
// The handler should check response_body_status() after this method returns
162+
// or from the callback to decide if the request should be aborted due to
163+
// transport failures.
164+
virtual CallbackStatus PartialReplyWithFlushCallback(
165+
std::function<void()> callback) = 0;
166+
virtual BodyStatus response_body_status() { return BodyStatus::PENDING; }
167+
168+
// Request streaming is disabled by default
169+
virtual BodyStatus request_body_status() { return BodyStatus::COMPLETE; }
146170

147171
// Completes the response and sends any buffered response body
148172
// to the client. Headers will be generated and sent first if PartialReply()
149173
// has never be called.
150174
// Assumes 200 if status is not specified.
151-
// Once Reply() is called, the request object will be owned (and destructed)
175+
// Once Reply() is called, the request object will be owned and destructed
152176
// by the server runtime.
153177
virtual void ReplyWithStatus(HTTPStatusCode status) = 0;
154178
virtual void Reply() = 0;
155179

156180
// Aborts the current request forcibly.
181+
// Once Abort() is called, the request object will be owned and destructed
182+
// by the server runtime.
157183
virtual void Abort() = 0;
158184

159-
// To be defined: status enum
160-
161185
protected:
162186
ServerRequestInterface() = default;
163187

0 commit comments

Comments
 (0)