187 Commits

Author SHA1 Message Date
OMGeeky
e9c86ce157 fix test needing serde1 feature 2024-01-19 22:45:53 +01:00
OMGeeky
d4f579542d Add test file that creates a compiler error when hygiene is violated/stuff is used from outer scope 2024-01-19 20:19:00 +01:00
OMGeeky
5011dbe057 Fix renaming of import & remove some not needed stuff (from earlier) 2024-01-19 20:18:00 +01:00
OMGeeky
4aa90ee933 Fully qualify Some, Ok and Err when inside quoted area 2024-01-19 18:41:28 +01:00
OMGeeky
eb91000fed cleanup std usages
see https://doc.rust-lang.org/reference/procedural-macros.html#procedural-macro-hygiene
2024-01-13 13:19:32 -08:00
Boris Sabatier
006a9f3af1 Fix cfg flag for tcp tests 2024-01-09 09:47:38 -08:00
Tim Kuehn
a3d00b07da Update imports for serde_transport unit tests 2024-01-02 13:45:37 -08:00
Boris Sabatier
d62706e62c feat: Allow to create tarpc on existing TCP listener 2024-01-02 13:43:27 -08:00
Boris Sabatier
b92dd154bc feat: Allow to create tarpc on existing UDS listener 2024-01-02 13:43:27 -08:00
Tim Kuehn
a6758fd1f9 BeforeRequest hook chaining.
It's unintuitive that serve.before(hook1).before(hook2) executes in
reverse order, with hook2 going before hook1. With BeforeRequestList,
users can write `before().then(hook1).then(hook2).serving(serve)`, and
it will run hook1, then hook2, then the service fn.
2023-12-29 23:13:06 -08:00
Tim Kuehn
2c241cc809 Fix readme and release notes. 2023-12-29 21:03:09 -08:00
Tim Kuehn
263ef8a897 Prepare for v0.34.0 release 2023-12-29 20:52:37 -08:00
Kevin Ge
d50290a21c Add README for example app 2023-12-29 20:32:00 -08:00
Kevin Ge
26988cb833 Update OpenTelemetry packages in example app 2023-12-29 20:32:00 -08:00
Tim Kuehn
6cf18a1caf Rewrite traits to use async-fn-in-trait.
- Stub
- BeforeRequest
- AfterRequest

Also removed the last remaining usage of an unstable feature,
iter_intersperse.
2023-12-29 13:52:05 -08:00
Tim Kuehn
84932df9b4 Return span in InFlightRequests::complete_request.
Rather than returning a bool, return the Span associated with the
request. This gives RequestDispatch more flexibility to annotate the
request span.
2023-12-29 13:52:05 -08:00
Tim Kuehn
8dc3711a80 Use async fn in generated traits!!
The major breaking change is that Channel::execute no longer internally
spawns RPC handlers, because it is no longer possible to place a Send
bound on the return type of Serve::serve. Instead, Channel::execute
returns a stream of RPC handler futures.

Service authors can reproduce the old behavior by spawning each response
handler (the compiler knows whether or not the futures can be spawned;
it's just that the bounds can't be expressed generically):

    channel.execute(server.serve())
           .for_each(|rpc| { tokio::spawn(rpc); })
2023-12-29 13:52:05 -08:00
Tim Kuehn
7c5afa97bb Add request hooks to the Serve trait.
This allows plugging in horizontal functionality, such as authorization,
throttling, or latency recording, that should run before and/or after
execution of every request, regardless of the request type.

The tracing example is updated to show off both client stubs as well as
server hooks.

As part of this change, there were some changes to the Serve trait:

1. Serve's output type is now a Result<Response, ServerError>..
   Serve previously did not allow returning ServerErrors, which
   prevented using Serve for horizontal functionality like throttling or
   auth. Now, Serve's output type is Result<Resp, ServerError>, making
   Serve a more natural integration point for horizontal capabilities.
2. Serve's generic Request type changed to an associated type. The
   primary benefit of the generic type is that it allows one type to
   impl a trait multiple times (for example, u64 impls TryFrom<usize>,
   TryFrom<u128>, etc.). In the case of Serve impls, while it is
   theoretically possible to contrive a type that could serve multiple
   request types, in practice I don't expect that to be needed.  Most
   users will use the Serve impl generated by #[tarpc::service], which
   only ever serves one type of request.
2023-12-29 13:52:05 -08:00
Tim Kuehn
324df5cd15 Add back the Client trait, renamed Stub.
Also adds a Client stub trait alias for each generated service.

Now that generic associated types are stable, it's almost possible to
define a trait for Channel that works with async fns on stable. `impl
trait in type aliases` is still necessary (and unstable), but we're
getting closer.

As a proof of concept, three more implementations of Stub are implemented;

1. A load balancer that round-robins requests between different stubs.
2. A load balancer that selects a stub based on a request hash, so that
   the same requests go to the same stubs.
3. A stub that retries requests based on a configurable policy.

   The "serde/rc" feature is added to the "full" feature because the Retry
   stub wraps the request in an Arc, so that the request is reusable for
   multiple calls.

   Server implementors commonly need to operate generically across all
   services or request types. For example, a server throttler may want to
   return errors telling clients to back off, which is not specific to any
   one service.
2023-12-29 13:52:05 -08:00
Guillaume Charmetant
3264979993 Fix warnings in README's example 2023-11-16 09:54:54 -08:00
Guillaume Charmetant
dd63fb59bf Fix tokio dep in the README's example
Add missing tokio feature in the example's dependencies.
2023-11-16 09:54:54 -08:00
Tim Kuehn
f4db8cc5b4 Address clippy lints 2023-11-16 00:00:27 -08:00
Tim Kuehn
e9ba350496 Update must-use UI tests 2023-11-16 00:00:27 -08:00
Tim Kuehn
e6d779e70b Remove mipsel workflow.
Mipsel was downgraded to tier 3, which broke this workflow.
https://github.com/rust-lang/compiler-team/issues/648
2023-11-16 00:00:27 -08:00
Izumi Raine
ce5f8cfb0c Simplify TLS example (#404) 2023-04-14 12:33:22 -07:00
Tim Kuehn
4b69dc8db5 Prepare release of v0.33.0 2023-04-03 11:03:55 -07:00
Bruno
866db2a2cd Bump opentelemetry to 0.18 (#401) 2023-04-03 10:38:17 -07:00
Tim Kuehn
bed85e2827 Prepare release of v0.32.0 2023-03-24 15:04:06 -07:00
Bruno
93f3880025 Return transport errors to the caller (#399)
* Make client::InFlightRequests generic over result.

Previously, InFlightRequests required the client response type to be a
server response. However, this prevented injection of non-server
responses: for example, if the client fails to send a request, it should
complete the request with an IO error rather than a server error.

* Gracefully handle client-side send errors.

Previously, a client channel would immediately disconnect when
encountering an error in Transport::try_send. One kind of error that can
occur in try_send is message validation, e.g. validating a message is
not larger than a configured frame size. The problem with shutting down
the client immediately is that debuggability suffers: it can be hard to
understand what caused the client to fail. Also, these errors are not
always fatal, as with frame size limits, so complete shutdown was
extreme.

By bubbling up errors, it's now possible for the caller to
programmatically handle them. For example, the error could be walked
via anyhow::Error:

```
    2023-01-10T02:49:32.528939Z  WARN client: the client failed to send the request

    Caused by:
        0: could not write to the transport
        1: frame size too big
```

* Some follow-up work: right now, read errors will bubble up to all pending RPCs. However, on the write side, only `start_send` bubbles up. `poll_ready`, `poll_flush`, and `poll_close` do not propagate back to pending RPCs. This is probably okay in most circumstances, because fatal write errors likely coincide with fatal read errors, which *do* propagate back to clients. But it might still be worth unifying this logic.

---------

Co-authored-by: Tim Kuehn <tikue@google.com>
2023-03-24 14:31:25 -07:00
cguentherTUChemnitz
878f594d5b Feature/tls over tcp example (#398)
Example tarpc service that encodes messages with bincode written to a TLS-over-TCP transport.

Certs were generated with openssl 3 using https://github.com/rustls/rustls/tree/main/test-ca

New dependencies:
- tokio-rustls to set up the TLS connections
- rustls-pemfile to load certs from .pem files
2023-03-22 10:35:21 -07:00
Tim Kuehn
aa9bbad109 Fix compile_fail tests for formatting changes on stable 2023-03-17 10:07:54 -07:00
Tim Kuehn
7e872ce925 Remove bad mem::forget usage.
mem::forget is a dangerous tool, and it was being used carelessly for
things that have safer alternatives. There was at least one bug where a
cloned tokio::sync::mpsc::UnboundedSender used for request cancellation
was being leaked on every successful server response, so its refcounts
were never decremented. Because these are atomic refcounts, they'll wrap
around rather than overflow when reaching the maximum value, so I don't
believe this could lead to panics or unsoundness.
2022-11-23 18:01:12 -08:00
Tim Kuehn
62541b709d Replace actions-rs 2022-11-23 16:39:29 -08:00
Tim Kuehn
8c43f94fb6 Remove unused Sealed trait 2022-11-17 00:57:56 -08:00
Tim Kuehn
7fa4e5064d Ignore clippy false positive 2022-11-13 00:25:07 -08:00
Tim Kuehn
94db7610bb Require a static lifetime for request_name. 2022-11-05 11:43:03 -07:00
Tim Kuehn
0c08d5e8ca Prepare release of v0.31.0 2022-11-03 13:29:46 -07:00
Tim Kuehn
75b15fe2aa Address clippy lint 2022-10-07 10:51:45 -07:00
Tim Kuehn
863a08d87e In example-service, print the port the server is listened on.
This is helpful when passing starting the server with --port 0.
2022-10-06 20:58:54 -07:00
Tim Kuehn
49ba8f8b1b Zero-pad the random number suffix of TempPathBufs.
This way, the hex number is always 16 digits, which is helpful for test
verification as well as simple consistency.
2022-10-03 18:50:50 -07:00
Kevin K
d832209da3 feat: Unix domain sockets with serde transports (#380)
* adds support for Unix Domain Socket generic transports
* adds a TempPathBuf that lives in temp and is removed on drop
2022-10-03 18:07:29 -07:00
royrustdev
584426d414 fix clippy warnings #378 2022-09-19 23:26:07 -07:00
royrustdev
50eb80c883 reference latest tarpc version in readme 2022-09-19 21:58:21 -07:00
royrustdev
1f0c80d8c9 bump github actions 2022-09-15 11:17:58 -07:00
Tim Kuehn
99bf3e62a3 Prepare release of 0.30.0 2022-08-12 16:08:33 -07:00
Tim Kuehn
68863e3db0 Remove Channel::request_cancellation.
This trait fn returns a private type, which means it's useless for
anyone using the Channel.

Instead, add an inert (now-public) ResponseGuard to TrackedRequest that,
when taken out of the ManuallyDrop, ensures a Channel's request state is
cleaned up. It's preferable to make ResponseGuard public instead of
RequestCancellations because it's a smaller API surface (no public
methods, just a Drop fn) and harder to misuse, because it is already
associated with the correct request ID to cancel.
2022-08-12 16:08:33 -07:00
Tim Kuehn
453ba1c074 Lower log level of log in the RPC callpath 2022-08-12 09:04:47 -07:00
Makro
e3eac1b4f5 Add LICENSE files to crates (#372) 2022-08-10 17:11:50 -07:00
kkharji
0e102288a5 feat: re-export used packages (#371)
## Problem
Library users might get stuck with or ran into issues while using tarpc because of incompatible third party libraries. in particular, tokio_serde and tokio_util.

## Solution
This PR does the following:

1. re-export tokio_serde as part of feature serde-transport, because the end user imports it to use some serde-transport APIs.
2. Update third library packages to latest release and fix resulting issues from that.

## Important Notes
tokio_util 7.3 DelayQueue::poll_expired API changed [0] therefore, InFlightRequests::poll_expired now returns Poll<Option<u64>>

[0] https://docs.rs/tokio-util/latest/tokio_util/time/delay_queue/struct.DelayQueue.html#method.poll_expired
2022-07-15 10:14:49 -07:00
Tim Kuehn
4c8ba41b2f #[allow(unstable_name_collisions)] for .ready()
.ready() is being added to std, but in the meantime, I don't want to stop using PollTest::ready.
2022-06-07 01:29:14 -07:00
Tim Kuehn
946c627579 Remove unused field 2022-06-07 01:29:14 -07:00
Tim Kuehn
104dd71bba Clean up Channel request data more reliably.
When an InFlightRequest is dropped before response completion, request
data in the Channel persists until either the request expires or the
client cancels the request. In rare cases, requests with very large
deadlines could clog up the Channel long after request processing
ceases.

This commit adds a drop hook to InFlightRequest so that if it is dropped
before execution completes, a cancellation message is sent to the
Channel so that it can clean up the associated request data.

This only works for when using `InFlightRequest::execute` or
`Channel::execute`. However, users of raw `Channel` have access
to the `RequestCancellation` handle via `Channel::request_cancellation`,
so they can implement a similar method if they wish to manually clean up
request data.

Note that once a Channel's request data is cleaned up, that request can
never be responded to, even if a response is produced afterward.

Fixes https://github.com/google/tarpc/issues/314
2022-06-07 01:29:04 -07:00
Tim Kuehn
012c481861 Move cancellation types into a dedicated module.
Cancellation utilities could be useful for both client and server code.
2022-06-05 18:54:52 -07:00
Tim Kuehn
dc12bd09aa Annotate types that impl Future with #[must_use].
These types do nothing unless polled / .awaited.
Annotating them with #[must_use] helps prevent a common class of coding errors.

Fixes https://github.com/google/tarpc/issues/368.
2022-06-05 18:54:52 -07:00
Tim Kuehn
2594ea8ce9 Prepare release of 0.29.0 2022-06-05 15:26:33 -07:00
Tim Kuehn
839b87e394 Serialize RPC deadline as a Duration.
Duration was previously serialized as SystemTime. However, absolute
times run into problems with clock skew: if the remote machine's clock
is too far in the future, the RPC deadline will be exceeded before
request processing can begin. Conversely, if the remote machine's clock
is too far in the past, the RPC deadline will not be enforced.

By converting the absolute deadline to a relative duration, clock skew
is no longer relevant, as the remote machine will convert the deadline
into a time relative to its own clock. This mirrors how the gRPC HTTP2
protocol includes a Timeout in the request headers [0] but the SDK uses
timestamps [1]. Keeping the absolute time in the core APIs maintains all
the benefits of today, namely, natural deadline propagation between RPC
hops when using the current context.

This serialization strategy means that, generally, the remote machine's
deadline will be slightly in the future compared to the local machine.
Depending on network transfer latencies, this could be microseconds to
milliseconds, or worse in the worst case. Because the deadline is not
intended for high-precision scenarios, I don't view this is as
problematic.

Because this change only affects the serialization layer, local
transports that bypass serialization are not affected.

[0] https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md
[1] https://grpc.io/blog/deadlines/#setting-a-deadline
2022-05-26 15:18:49 -07:00
Tim Kuehn
57d0638a99 Add rpc.deadline tag to Opentelemetry traces. 2022-05-26 15:18:49 -07:00
Tim Kuehn
a3a6404a30 Prepare release of 0.28.0 2022-04-06 22:07:07 -07:00
Tim Kuehn
b36eac80b1 Bump minimum rust version to 1.58.0 2022-04-06 21:53:56 -07:00
Bruno
d7070e4bc3 Update opentelemetry and related dependencies (#362) 2022-04-03 14:09:14 -07:00
Tim Kuehn
b5d1828308 Use captured identifiers in format strings.
This was stabilized in Rust 1.58.0: https://blog.rust-lang.org/2022/01/13/Rust-1.58.0.html
2022-01-13 15:00:44 -08:00
Zak Cutner
92cfe63c4f Use single-threaded Tokio runtime (#360) 2022-01-06 20:59:51 -08:00
Tim Kuehn
839a2f067c Update example to latest version of Clap 2021-12-27 22:56:17 -08:00
David Kleingeld
b5d593488c Derive more traits for RpcError (#359)
Makes RpcError derive Clone, PartialEq, Eq, Hash, Serialize and Deserialize.
2021-12-27 22:00:58 -08:00
Tim Kuehn
eea38b8bf4 Simplify code with const assert!.
The code that prevents compilation on systems where usize is larger than
u64 previously used a const index-out-of-bounds trick. That code can now
be replaced with assert!, as const panic! has landed in 1.57.0 stable.
2021-12-03 15:20:33 -08:00
Shi Yan
70493c15f4 Fix a compiling issue of the official example (#358)
Fix a compiling issue of the official example because of the following error :

```
error[E0599]: the method `execute` exists for struct `BaseChannel<_, _, UnboundedChannel<ClientMessage<_>, Response<_>>>`, but its trait bounds were not satisfied
  --> src/main.rs:39:25
   |
39 |     tokio::spawn(server.execute(HelloServer.serve()));
   |                         ^^^^^^^ method cannot be called on `BaseChannel<_, _, UnboundedChannel<ClientMessage<_>, Response<_>>>` due to unsatisfied trait bounds
   |
   = note: the following trait bounds were not satisfied:
           `<&BaseChannel<_, _, UnboundedChannel<ClientMessage<_>, Response<_>>> as futures::Stream>::Item = _`
           which is required by `&BaseChannel<_, _, UnboundedChannel<ClientMessage<_>, Response<_>>>: tarpc::server::incoming::Incoming<_>`
           `&BaseChannel<_, _, UnboundedChannel<ClientMessage<_>, Response<_>>>: futures::Stream`
           which is required by `&BaseChannel<_, _, UnboundedChannel<ClientMessage<_>, Response<_>>>: tarpc::server::incoming::Incoming<_>`
   = help: items from traits can only be used if the trait is in scope
help: the following trait is implemented but not in scope; perhaps add a `use` for it:
   |
1  | use tarpc::server::Channel;
   |
```

See https://github.com/google/tarpc/pull/358#issuecomment-981953193 for the root cause.
2021-11-29 17:01:16 -08:00
baptiste0928
f7c5d6a7c3 Fix example-service (#355)
Fixes the compilation of the example-service crate (the Clap trait has been renamed Parser in clap-rs/clap@d840d56).
2021-11-15 08:39:04 -08:00
Scott Kirkpatrick
98c5d2a18b Re-add typo fixes (#353)
The typo fixes that were added by commit b5d9aaa
were accidentally reverted by commit 1e680e3, this
will add them back
2021-11-08 10:07:21 -08:00
Tim Kuehn
46b534f7c6 Use HashMap::shrink_to in impl of Comapct::compact. 2021-10-21 17:03:57 -07:00
Tim Kuehn
42b4fc52b1 Set rust-version to 1.56 2021-10-21 16:08:15 -07:00
Tim Kuehn
350dbcdad0 Upgrade to Rust 2021! 2021-10-21 14:10:21 -07:00
Tim Kuehn
b1b4461d89 Prepare release of 0.27.2 2021-10-08 22:31:56 -07:00
Tim Kuehn
f694b7573a Close TcpStream when client disconnects.
An attempt at a clean shutdown helps the server to drop its connections
more quickly.

Testing this uncovered a latent bug in DelayQueue wherein `poll_expired`
yields `Pending` when empty. A workaround was added to
`InFlightRequests::poll_expired`: check if there are actually any
outstanding requests before calling `DelayQueue::poll_expired`.
2021-10-08 22:13:24 -07:00
Tim Kuehn
1e680e3a5a Fix typos in docs.
Fixes https://github.com/google/tarpc/issues/352.
2021-10-08 19:19:50 -07:00
Tim Kuehn
2591d21e94 Update release notes to mention io::Error = 2021-09-23 13:57:43 -07:00
Tim Kuehn
6632f68d95 Prepare for 0.27 release 2021-09-22 15:41:34 -07:00
Dmitry Kakurin
25985ad56a Update README.md (#350)
Fixed 2 typos
2021-09-01 17:58:49 -07:00
Tim Kuehn
d6a24e9420 Address Clippy lint 2021-08-24 12:40:18 -07:00
Tim Kuehn
281a78f3c7 Add tokio-serde-bincode feature 2021-08-24 12:37:57 -07:00
Julian Tescher
a0787d0091 Update to opentelemetry 0.16.x (#349) 2021-08-17 00:00:07 -04:00
Frederik-Baetens
d2acba0e8a add serde-transport-json feature flag (#346)
In general, it should be possible to use, or at least import all functionality of a library, when having only that library in your cargo.toml.
2021-05-06 08:41:57 -07:00
Tim Kuehn
ea7b6763c4 Refactor server module.
In the interest of the user's attention, some ancillary APIs have been
moved to new submodules:

- server::limits contains what was previously called Throttler and
  ChannelFilter. Both of those names were very generic, when the methods
  applied by these types were very specific (and also simplistic). Renames
  have occurred:
  - ThrottlerStream => MaxRequestsPerChannel
  - Throttler => MaxRequests
  - ChannelFilter => MaxChannelsPerKey
- server::incoming contains the Incoming trait.
- server::tokio contains the tokio-specific helper types.

The 5 structs and 1 enum remaining in the base server module are all
core to the functioning of the server.
2021-04-21 17:05:49 -07:00
Tim Kuehn
eb67c540b9 Use more structured errors in client. 2021-04-21 14:54:45 -07:00
Tim Kuehn
4151d0abd3 Move Span creation into BaseChannel.
It's important for Channel decorators, like Throttler, to have access to
the Span. This means that the BaseChannel becomes responsible for
starting its own requests. Actually, this simplifies the integration for
the Channel users, as they can assume any yielded requests are already
tracked.

This entails the following breaking changes:

- removed trait method Channel::start_request as it is now done
  internally.
2021-04-21 14:54:45 -07:00
Tim Kuehn
d0c11a6efa Change RPC error type from io::Error => RpcError.
Becaue tarpc is a library, not an application, it should strive to
use structured errors in its API so that users have maximal flexibility
in how they handle errors. io::Error makes that hard, because it is a
kitchen-sink error type.

RPCs in particular only have 3 classes of errors:

- The connection breaks.
- The request expires.
- The server decides not to process the request.

(Of course, RPCs themselves can have application-specific errors, but
from the perspective of the RPC library, those can be classified as
successful responsees).
2021-04-20 18:29:55 -07:00
Tim Kuehn
82c4da1743 Prepare release of v0.26.2 2021-04-20 11:28:15 -07:00
Tim Kuehn
0a15e0b75c Rustdoc: link RPC futures to their methods. 2021-04-20 11:25:26 -07:00
Tim Kuehn
0b315c29bf It's not currently possible to document the enum variants, which means
projects that #[deny(missing_docs)] wouldn't compile if using tarpc
services.
2021-04-20 09:01:39 -07:00
Tim Kuehn
56f09bf61f Fix log that's split across lines. 2021-04-17 17:15:16 -07:00
Tim Kuehn
6d82e82419 Fix formatting 2021-04-16 16:51:21 -07:00
Tim Kuehn
9bebaf814a Address clippy lint 2021-04-14 17:49:27 -07:00
Tim Kuehn
5f4d6e6008 Prepare release of v0.26.0 2021-04-14 17:08:44 -07:00
Tim Kuehn
07d07d7ba3 Remove tracing_appender as it does not support build target mipsel-unknown-linux-gnu 2021-04-01 19:37:02 -07:00
Tim Kuehn
a41bbf65b2 Use rustfmt instead of cargo fmt so that diff is only printed once 2021-04-01 17:24:34 -07:00
Tim Kuehn
21e2f7ca62 Tear out requirement that Transport's error type is io::Error. 2021-04-01 17:24:34 -07:00
Tim Kuehn
7b7c182411 Instrument tarpc with tracing.
tarpc is now instrumented with tracing primitives extended with
OpenTelemetry traces. Using a compatible tracing-opentelemetry
subscriber like Jaeger, each RPC can be traced through the client,
server, amd other dependencies downstream of the server. Even for
applications not connected to a distributed tracing collector, the
instrumentation can also be ingested by regular loggers like env_logger.

 # Breaking Changes

 ## Logging

Logged events are now structured using tracing. For applications using a
logger and not a tracing subscriber, these logs may look different or
contain information in a less consumable manner. The easiest solution is
to add a tracing subscriber that logs to stdout, such as
tracing_subscriber::fmt.

 ##  Context

- Context no longer has parent_span, which was actually never needed,
  because the context sent in an RPC is inherently the parent context.
  For purposes of distributed tracing, the client side of the RPC has all
  necessary information to link the span to its parent; the server side
  need do nothing more than export the (trace ID, span ID) tuple.
- Context has a new field, SamplingDecision, which has two variants,
  Sampled and Unsampled. This field can be used by downstream systems to
  determine whether a trace needs to be exported. If the parent span is
  sampled, the expectation is that all child spans be exported, as well;
  to do otherwise could result in lossy traces being exported. Note that
  if an Openetelemetry tracing subscriber is not installed, the fallback
  context will still be used, but the Context's sampling decision will
  always be inherited by the parent Context's sampling decision.
- Context::scope has been removed. Context propagation is now done via
  tracing's task-local spans. Spans can be propagated across tasks via
  Span::in_scope. When a service receives a request, it attaches an
  Opentelemetry context to the local Span created before request handling,
  and this context contains the request deadline. This span-local deadline
  is retrieved by Context::current, but it cannot be modified so that
  future Context::current calls contain a different deadline. However, the
  deadline in the context passed into an RPC call will override it, so
  users can retrieve the current context and then modify the deadline
  field, as has been historically possible.
- Context propgation precedence changes: when an RPC is initiated, the
  current Span's Opentelemetry context takes precedence over the trace
  context passed into the RPC method. If there is no current Span, then
  the trace context argument is used as it has been historically. Note
  that Opentelemetry context propagation requires an Opentelemetry
  tracing subscriber to be installed.

 ## Server

- The server::Channel trait now has an additional required associated
  type and method which returns the underlying transport. This makes it
  more ergonomic for users to retrieve transport-specific information,
  like IP Address. BaseChannel implements Channel::transport by returning
  the underlying transport, and channel decorators like Throttler just
  delegate to the Channel::transport method of the wrapped channel.

 # References

[1] https://github.com/tokio-rs/tracing
[2] https://opentelemetry.io
[3] https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-jaeger
[4] https://github.com/env-logger-rs/env_logger
2021-04-01 17:24:34 -07:00
Ben Ludewig
db0c778ead Serialize u128 TraceId as LE bytes (#344) 2021-03-30 08:41:19 -07:00
Tim Kuehn
c3efb83ac1 Add more context to errors returned by serde transport 2021-03-28 20:03:03 -07:00
Tim Kuehn
3d7b0171fe Fix cargo fmt portion of pre-commit 2021-03-26 19:39:56 -07:00
oblique
c191ff5b2e Do not enable tokio-serde/json by default (#345) 2021-03-26 18:22:44 -07:00
Tim Kuehn
90bc7f741d Fix up imports 2021-03-17 12:44:39 -07:00
Kitsu
d3f6c01df2 Reduce required tokio features (#343)
* Move async tests behind cfg-ed mod
* Use explicit tokio features for the example
* Use only relative crate path for example dependency
2021-03-17 12:30:18 -07:00
Tim Kuehn
c6450521e6 Add method to run a future in the current context.
Previously, `Context::current` would always return a new context. Now,
it uses tokio task-local data to look for the current context. Tokio
task locals are not actually tied to a tokio executor; instead, they
provide data scoped to a future.

The basic pattern is:

```rust
let ctx = Context::new_root();
ctx.scope(async {
    let ctx2 = context::current();
    assert_eq!(ctx2.trace_context.span_id, ctx.trace_context.span_id);
});
```

`server::InFlightRequest::execute` uses `Context::scope` to set the
current context before executing a request, so calls to
`context::current` in request handlers will return the context provided
by the client. This does not propagate to new spawned tasks. To
propagate the client context to child tasks, the following pattern will
work:

```rust
tokio::spawn(context::current().scope(async { /* do work here */ }));
```

This commit also introduces a breaking change to Context serialization.
Previously, the deadline only serialized second-level precision. Now, it
provides full fidelity serialization to the nanosecond.
2021-03-13 16:05:02 -08:00
Tim Kuehn
1da6bcec57 Prepare v0.25 release 2021-03-10 20:00:25 -08:00
Seth Vargo
75a5591158 Improve Actions hygiene
👋 hello there! I'm a fellow Googler who works on projects that leverage GitHub Actions for CI/CD. Recently I noticed a large increase in our queue time, and I've tracked it down to the [limit of 180 concurrent jobs](https://docs.github.com/en/actions/reference/usage-limits-billing-and-administration) for an organization. To help be better citizens, I'm proposing changes across a few repositories that will reduce GitHub Actions hours and consumption. I hope these changes are reasonable and I'm happy to talk through them in more detail.

- Only run GitHub Actions for pushes and PRs against the main branch of the repository. If your team uses a forking model, this change will not affect you. If your team pushes branches to the repository directly, this changes actions to only run against the primary branches or if you open a Pull Request against a primary branch.

- For long-running jobs (especially tests), I added the "Cancel previous" workflow. This is very helpful to prevent a large queue backlog when you are doing rapid development and pushing multiple commits. Without this, GitHub Actions' default behavior is to run all actions on all commits.

There are other changes you could make, depending on your project (but I'm not an expert):

- If you have tests that should only run when a subset of code changes, consider gating your workflow to particular file paths. For example, we have some jobs that do Terraform linting, but [they only run when Terraform files are changed](c4f59fee71/.github/workflows/terraform.yml (L3-L11)).

Hopefully these changes are not too controversial and also hopefully you can see how this would reduce actions consumption to be good citizens to fellow Googlers. If you have any questions, feel free to respond here or ping me on chat. Thank you!
2021-03-10 17:31:13 -08:00
Tim Kuehn
9462aad3bf Improve test coverage of serde_transport 2021-03-10 12:37:55 -08:00
Tim Kuehn
0964fc51ff Add transport::channel::bounded.
This is like transport::channel::unbounded but with a fixed buffer size.
2021-03-08 23:10:12 -08:00
Tim Kuehn
27aacab432 Alternate polling expired and new requests.
Previously, there were two loops:

- Expired in-flight requests are polled until Pending.
- New requests are polled until Pending.

Now there is one loop that alternates between polling expired requests
and new requests. This way, neither type of action can face starvation.
2021-03-08 23:00:39 -08:00
Tim Kuehn
3feb465ad3 Clean up some server documentation. 2021-03-08 15:43:23 -08:00
Tim Kuehn
66cdc99ae0 Factor out ensure_writeable methods.
There is some important logic that is easy to overlook in the client and
server channels: streams of data to write to the transport should not be
polled until the transport is known to be ready to buffer a message. In
the case that a transport's buffer is full, it needs to be flushed to
make room for more messages.

Without this logic, start_send() could return an error when the buffer
is full, which would cause the entire Channel to error out.

Due to the importance of this logic, it's now factored out into its own
method that's easier to understand: fn ensure_writeable. There is one in
the client module and and one in the server module.
2021-03-08 11:36:20 -08:00
Tim Kuehn
66419db6fd Don't send a deadline-exceeded response.
The deadline-exceeded response was largely redundant, because the client
shouldn't normally be waiting for such a response, anyway -- the normal
client will automatically remove the in-flight request when it reaches
the deadline.

This also allows for internalizing the expiration+cleanup logic entirely
within BaseChannel, without having it leak into the Channel trait and
requiring action taken by the Requests struct.
2021-03-07 23:49:31 -08:00
Tim Kuehn
72d5dbba89 Cleanup wrap-up.
- Remove unnecessary Sync and Clone bounds.
- Merge client and client::channel modules.
- Run cargo clippy in the pre-push hook.
- Put DispatchResponse.cancellation in an Option.  Previously, the
  cancellation logic looked to see if `complete == true`, but it's a bit
  less error prone to put the Cancellation in an Option, so that the
  request can't accidentally be cancelled.
- Remove some unnecessary pins/projections.
- Clean up docs a bit. rustdoc had some warnings that are now gone.
2021-03-07 22:29:03 -08:00
Tim Kuehn
e75193c191 Client RPCs now take &self.
This required the breaking change of removing the Client trait. The
intent of the Client trait was to facilitate the decorator pattern by
allowing users to create their own Clients that added behavior on top of
the base client. Unfortunately, this trait had become a maintenance
burden, consistently causing issues with lifetimes and the lack of
generic associated types. Specifically, it meant that Client impls could
not use async fns, which is no longer tenable today.
2021-03-07 17:41:29 -08:00
Tim Kuehn
ce4fd49161 Centralize client-side request deadline handling.
Before this commit, each request future had its own timeout and would
communicate to the client Channel when a request was no longer being
listened to. Now, instead, the Channel tracks deadlines of in-flight
requests and completes requests with deadline-exceeded errors when they
expire.

This should be functionally equivalent to the previous way. It just cuts
down on the amount of two-way processing required. Unfortunately,
dropping a response future early still requires the client to send a
cancellation message to the Channel.
2021-03-07 15:31:17 -08:00
Tim Kuehn
3c978c5bf6 Handle deadlines in BaseChannel.
Before this commit, deadlines were handled by a timeout future that
wrapped each request handler. However, request handlers can be dropped
before sending a response back to the channel, so they can't be relied
on for channel state cleanup. Additionally, clients can't be relied on
to send cancellation messages. It was therefore theoretically possible
for pathological behaviors to cause an unbounded growth in orphan
request data in the Channel.

With this change, as long as requests sent have reasonable deadlines,
then the channel will be able to clean itself up. It is still possible
for requests to be sent with very large deadlines, which would prevent
the channel from cleaning itself up.
2021-03-07 04:05:33 -08:00
Tim Kuehn
6f419e9a9a Refactor server module to be easier to understand.
1. Renames

Some of the items in this module were renamed to be less generic:

- Handler => Incoming
- ClientHandler => Requests
- ResponseHandler => InFlightRequest
- Channel::{respond_with => requests}

In the case of Handler: handler of *what*? Now it's a bit clearer that
this is a stream of Channels (aka *incoming* connections).

Similarly, ClientHandler was a stream of requests over a single
connection. Hopefully Requests better reflects that.

ResponseHandler was renamed InFlightRequest because it no longer
contains the serving function. Instead, it is just the request, plus
the response channel and an abort hook. As a result of this,
Channel::respond_with underwent a big change: it used to take the
serving function and return a ClientHandler; now it has been renamed
Channel::requests and does not take any args.

2. Execute methods

All methods thats actually result in responses being generated
have been consolidated into methods named `execute`:

- InFlightRequest::execute returns a future that completes when a
  response has been generated and sent to the server Channel.
- Requests::execute automatically spawns response handlers for all
  requests over a single channel.
- Channel::execute is a convenience for `channel.requests().execute()`.
- Incoming::execute automatically spawns response handlers for all
  requests over all channels.

3. Removal of Server.

server::Server was removed, as it provided no value over the Incoming/Channel
abstractions. Additionally, server::new was removed, since it just
returned a Server.
2021-03-06 20:20:48 -08:00
Tim Kuehn
b3eb8d0b7a Move items in the rpc module to the top level.
The rpc module doesn't carry its weight. The whole darn project is RPC related!
2021-03-06 15:05:10 -08:00
Tim Kuehn
3b422eb179 Abort all in-flight requests when dropping BaseChannel.
Fixes #341
2021-01-24 17:57:44 -08:00
Michael Zimmermann
4b513bad73 fix clippy::needless_lifetimes
warning: explicit lifetimes given in parameter types where they could be elided (or replaced with `'_` if needed by type declaration)
   --> tarpc/src/rpc/server/filter.rs:127:5
    |
127 |     fn channel<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut C> {
    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
    = note: `#[warn(clippy::needless_lifetimes)]` on by default
    = help: for further information visit https://rust-lang.github.io/rust-clippy/master/index.html#needless_lifetimes

warning: 1 warning emitted
2021-01-20 23:27:50 -08:00
Michael Zimmermann
e71e17866d github actions: cargo-check mipsel-unknown-linux-gnu 2021-01-20 23:27:50 -08:00
Michael Zimmermann
7e3fbec077 example-service: set max frame length to usize::MAX
I don't know what the intention was behind using u32::MAX + 1 but since the
argument's type is usize this is the only giant value that makes sense to me.
2021-01-20 23:27:50 -08:00
Michael Zimmermann
e4bc5e8e32 use AtomicUsize instead of AtomicU64
- it's more portable (some architectures like MIPS don't support AtomicU64)
- for most 64bit architectures usize should be 64bit as well
- for most users even 32bit would probably be enough because:
  - it's tied to the connection(for streaming sockets)
  - the ID wraps and by the time that happens, all previous requests would have
    timed out unless you send a lot of requests and have a ton of RAM
2021-01-20 23:27:50 -08:00
Tim Kuehn
bc982c5584 Prepare release of v0.24.1 2020-12-28 15:42:11 -08:00
Logan Magee
d440e12c19 Bump tokio to 1.0 (#337)
Co-authored-by: Artem Vorotnikov <artem@vorotnikov.me>
2020-12-23 22:49:02 -08:00
Frederik-Baetens
bc8128af69 add serde derivation alias macro (#333) 2020-11-13 14:36:59 -08:00
Tim Kuehn
1d87c14262 Fix github actions config - take 3 2020-11-12 12:33:10 -08:00
Tim Kuehn
ca929c2178 Fix github actions config - take 2 2020-11-12 12:24:46 -08:00
Tim Kuehn
569039734b Fix github actions config 2020-11-12 12:13:10 -08:00
Tim Kuehn
3d43310e6a Make 'cargo test' succeed again 2020-11-12 11:59:39 -08:00
Tim Kuehn
d21cbddb0d Cargo test should pass without features enabled 2020-11-12 11:57:08 -08:00
Frederik-Baetens
25aa857edf Reexport/tokio serde (#332)
Re-export tokio_serde when the serde-transport feature is enabled.
2020-11-09 12:56:46 -08:00
Frederik-Baetens
0bb2e2bbbe re-export serde (#330)
* re-export serde
* make serde re-export dependent on serde1 feature flag
* update missing_async compile test case
2020-11-09 11:42:28 -08:00
chansuke
dc376343d6 Remove #[derive(Debug)] from library structs (#327)
* Remove `#[derive(Debug)]` from library structs
* Add manual debug impl for backward compatibility
2020-11-04 11:24:57 -08:00
Artem Vorotnikov
2e7d1f8a88 Bump dependencies (#328) 2020-10-31 09:43:40 -07:00
Tim Kuehn
6314591c65 Add tokio's macros feature to readme example's dependencies 2020-10-30 17:29:14 -07:00
Tim Kuehn
7dd7494420 Prepare v0.23.1 release 2020-10-29 18:54:35 -07:00
Tim Kuehn
6c10e3649f Fix tokio required features 2020-10-29 18:53:04 -07:00
Tim Kuehn
4c6dee13d2 cargo fmt 2020-10-29 00:44:15 -07:00
Bernardo Meurer
e45abe953a tarpc: enable tokio's time feature (#325) 2020-10-29 00:43:38 -07:00
Tim Kuehn
dec3e491b5 Fix unused import 2020-10-27 15:52:11 -07:00
Kitsu
6ce341cf79 Add example for custom transport usage (#322) 2020-10-23 14:28:26 -07:00
Tim Kuehn
b9868250f8 Prepare release of v0.23.0 2020-10-19 11:12:43 -07:00
Urhengulas
a3f1064efe Cargo.toml: Clean + update dependencies 2020-10-18 16:03:04 -07:00
Johann Hemmann
026083d653 Bump tokio from 0.2 to 0.3 (#319)
# Bump `tokio` from 0.2 to 0.3

* `Cargo.toml`:
    * bump `tokio` from 0.2 to 0.3
    * bump `tokio-util` from 0.3 to 0.4
    * remove feature `time` from `tokio`
    * fix alphabetical order of dependencies
* `tarpc::rpc`:
    * `client, server`: `tokio::time::Elapsed` -> `tokio::time::error::Elapsed`
    * `client, transport`, `::tests`: Fix `#[tokio::test]` macro usage
* `tarpc::serde_transport`:
    * `TcpListener.incoming().poll_next(...)` -> `TcpListener.poll_accept(...)`
      -> https://github.com/tokio-rs/tokio/discussions/2983
    * Adapt `AsyncRead`, `AsynWrite` implements in tests
* `README.md`, `tarpc::lib`: Adapt tokio version in docs

# Satisfy clippy

* replace `match`-statements with `matches!(...)`-macro
2020-10-17 17:33:08 -07:00
Tim Kuehn
d27f341bde Prepare release of v0.22.0 2020-08-19 18:35:36 -07:00
Tim Kuehn
2264ebecfc Remove serde_transport::tcp::connect_with.
Instead, serde_transport::tcp::connect returns a future named Connect
that has methods to directly access the framing config. This is
consistent with how serde_transport::tcp::listen returns a future with
methods to access the framing config. In addition to this consistency,
it reduces the API surface and provides a simpler user transition from
"zero config" to "some config".
2020-08-19 17:51:53 -07:00
Tim Kuehn
3207affb4a Update pre-commit for changes to cargo fmt.
--write-mode is now --check.
2020-08-19 17:51:20 -07:00
Andre B. Reis
0602afd50c Make connect() and connect_with() take a FnOnce for the codec (#315) 2020-08-19 16:15:26 -07:00
Tim Kuehn
4343e12217 Fix incorrect documentation 2020-08-18 02:58:11 -07:00
Tim Kuehn
7fda862fb8 Run cargo fmt 2020-08-18 02:55:24 -07:00
Tim Kuehn
aa7b875b1a Expose framing config in serde_transport. 2020-08-18 02:47:41 -07:00
Tim Kuehn
54d6e0e3b6 Add license headers 2020-08-04 17:33:41 -07:00
Tim Kuehn
bea3b442aa Move mod.rs files up one directory.
It's easier in IDEs if the files aren't all named the same.
2020-08-04 17:25:53 -07:00
Tim Kuehn
954a2502e7 Remove duplicate rustdoc 2020-08-02 22:24:09 -07:00
Tim Kuehn
e3f34917c5 Prepare v0.21.1 2020-08-02 21:34:13 -07:00
Tim Kuehn
f65dd05949 Enable documentation for optional features on docs.rs 2020-08-02 20:57:21 -07:00
Tim Kuehn
240c436b34 Ensure Context is Sync. 2020-08-01 14:01:07 -07:00
Tim Kuehn
c9803688cc Ensure Context is Send. 2020-08-01 13:49:25 -07:00
Tim Kuehn
4987094483 Compression example.
Follow-up work: some extension points would be useful allow enabling compression on a per-request basis.

Fixes https://github.com/google/tarpc/issues/200
2020-08-01 13:45:16 -07:00
Tim Kuehn
ff55080193 Minor refactor 2020-07-30 13:11:13 -07:00
Tim Kuehn
258193c932 PubSub example needs to populate the subscription topics. 2020-07-30 11:14:13 -07:00
Tim Kuehn
67823ef5de Get rid of sleeps in PubSub example. 2020-07-30 01:27:31 -07:00
Tim Kuehn
a671457243 Add topics to PubSub example 2020-07-29 22:51:04 -07:00
Tim Kuehn
cf654549da Add documentation to PubSub example. 2020-07-29 18:05:35 -07:00
Tim Kuehn
6a01e32a2d Shut down client dispatch immediately when read half of transport is closed.
Clients can't receive any responses when the read half is closed, which means they can't verify if their requests were served. Therefore, there is no point in writing further requests after the read half is closed.
2020-07-29 13:50:42 -07:00
Tim Kuehn
e6597fab03 Add some error context to client dispatch.
I'm taking this opportunity to experiment with anyhow. So far, results are promising. It was a bit hard to use with Poll<Option<Result<T, E>>> types, so I added a crate-internal helper trait for that.
2020-07-29 12:07:07 -07:00
Tim Kuehn
ebd245a93d Rewrite pubsub example to have the subscriber connect to the publisher.
Fixes https://github.com/google/tarpc/issues/313
2020-07-28 22:10:17 -07:00
Tim Kuehn
3ebc3b5845 Add accessor fns.
- ClientHandler::get_pin_channel
- BaseChannel::get_pin_ref
- serde_transport::Transport::get_ref
2020-07-28 21:27:36 -07:00
Tim Kuehn
0e5973109d Make docs.rs document feature-gated public items. 2020-07-28 19:43:43 -07:00
Tim Kuehn
5f02d7383a Add tests for correct diagnostic output from proc macro-generated compiler errors. 2020-07-27 01:17:06 -07:00
Tim Kuehn
2bae148529 Address clippy lints 2020-07-27 00:04:45 -07:00
Tim Kuehn
42a2e03aab Add better diagnostics for missing 'async' in impls using #[tarpc::server] 2020-07-26 23:47:48 -07:00
Tim Kuehn
b566d0c646 Use #[tarpc::server] in example-service 2020-07-26 18:26:41 -07:00
Jon Cinque
b359f16767 Add concurrent tests using join and join_all
These tests are essentially copies of the `concurrent` test,
specifically using `join` and `join_all`.  Note that for the `join_all`
example to work, all of the `Client` clones must be created before *any*
requests are added, otherwise there will be a lifetime problem with the
second request, saying that second client, `c2`, is still borrowed when
`req1` is dropped.  It would require a larger redesign to fix this
issue.
2020-07-24 09:51:05 -07:00
Greg Fitzgerald
f8681ab134 Migrate examples to tarpc::server 2020-07-22 14:03:23 -07:00
Tim Kuehn
7e521768ab Prepare for v0.21.0 release. 2020-06-26 20:05:02 -07:00
Tim Kuehn
e9b1e7d101 Use #[non_exhaustive] in lieu of _NonExhaustive enum variant. 2020-06-26 19:47:20 -07:00
Taiki Endo
f0322fb892 Remove uses of pin_project::project attribute
pin-project will deprecate the project attribute due to some unfixable
limitations.

Refs: https://github.com/taiki-e/pin-project/issues/225
2020-06-05 20:34:44 -07:00
Patrick Elsen
617daebb88 Add tarpc::server proc-macro as syntactic sugar for async methods. (#302)
The tarpc::server proc-macro can be used to annotate implementations of
services to turn async functions into the proper declarations needed
for tarpc to be able to call them.

This uses the assert_type_eq crate to check that the transformations
applied by the tarpc::server proc macro are correct and lead to code
that compiles.
2020-05-16 10:25:25 -07:00
Tim Kuehn
a11d4fff58 Remove raii_counter 2020-04-22 02:13:02 -07:00
Tim
bf42a04d83 Move the request timeout so that it surrounds the entire call, not just the response future. (#295)
* Move the request timeout so that it surrounds the entire call, not just the response future.

This will enable the timeout earlier, so that a backlog in the outbound request buffer can not cause requests to stall indefinitely.

* Run cargo fmt
2020-02-25 14:42:40 -08:00
Tim Kuehn
06528d6953 Fix clippy lint. 2019-12-19 12:28:26 -08:00
Tim Kuehn
9f00395746 Replace _non_exhaustive fields with #[non_exhaustive] attribute.
The attribute landed on stable rust (1.40.0) today.

Fixes https://github.com/google/tarpc/issues/275
2019-12-19 12:14:34 -08:00
Tim Kuehn
e0674cd57f Make pre-push run on rust stable. 2019-12-19 12:06:06 -08:00
Tim Kuehn
7e49bd9ee7 Clean up badges a bit. 2019-12-16 13:21:00 -08:00
Tim Kuehn
8a1baa9c4e Remove usage of unsafe in rpc::client::channel.
pin_project is actually able to handle the complexities of enum Futures.
2019-12-16 11:10:57 -08:00
Oleg Nosov
31c713d188 Allow raw identifiers + fixed naming + place all code generation methods in impl (#291)
Allows defining services using raw identifiers like:

```rust
pub mod service {
    #[tarpc::service]
    pub trait r#trait {
        async fn r#fn(x: i32) -> Result<u8, String>;
    }
}
```

Also:

- Refactored names (ident -> type)
- All code generation methods placed in impl
2019-12-12 10:13:57 -08:00
80 changed files with 8574 additions and 3954 deletions

View File

@@ -1,4 +1,10 @@
on: [push, pull_request]
on:
push:
branches:
- master
pull_request:
branches:
- master
name: Continuous integration
@@ -7,60 +13,55 @@ jobs:
name: Check
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v1
- uses: actions-rs/toolchain@v1
- name: Cancel previous
uses: styfle/cancel-workflow-action@0.10.0
with:
profile: minimal
toolchain: stable
override: true
- uses: actions-rs/cargo@v1
with:
command: check
args: --all-features
access_token: ${{ github.token }}
- uses: actions/checkout@v3
- uses: dtolnay/rust-toolchain@stable
- run: cargo check --all-features
test:
name: Test Suite
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v1
- uses: actions-rs/toolchain@v1
- name: Cancel previous
uses: styfle/cancel-workflow-action@0.10.0
with:
profile: minimal
toolchain: stable
override: true
- uses: actions-rs/cargo@v1
with:
command: test
args: --all-features
access_token: ${{ github.token }}
- uses: actions/checkout@v3
- uses: dtolnay/rust-toolchain@stable
- run: cargo test
- run: cargo test --manifest-path tarpc/Cargo.toml --features serde1
- run: cargo test --manifest-path tarpc/Cargo.toml --features tokio1
- run: cargo test --manifest-path tarpc/Cargo.toml --features serde-transport
- run: cargo test --manifest-path tarpc/Cargo.toml --features tcp
- run: cargo test --all-features
fmt:
name: Rustfmt
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v1
- uses: actions-rs/toolchain@v1
- name: Cancel previous
uses: styfle/cancel-workflow-action@0.10.0
with:
profile: minimal
toolchain: stable
override: true
- run: rustup component add rustfmt
- uses: actions-rs/cargo@v1
access_token: ${{ github.token }}
- uses: actions/checkout@v3
- uses: dtolnay/rust-toolchain@stable
with:
command: fmt
args: --all -- --check
components: rustfmt
- run: cargo fmt --all -- --check
clippy:
name: Clippy
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v1
- uses: actions-rs/toolchain@v1
- name: Cancel previous
uses: styfle/cancel-workflow-action@0.10.0
with:
profile: minimal
toolchain: stable
override: true
- run: rustup component add clippy
- uses: actions-rs/cargo@v1
access_token: ${{ github.token }}
- uses: actions/checkout@v3
- uses: dtolnay/rust-toolchain@stable
with:
command: clippy
args: --all-features -- -D warnings
components: clippy
- run: cargo clippy --all-features -- -D warnings

View File

@@ -1,7 +1,11 @@
[workspace]
resolver = "2"
members = [
"example-service",
"tarpc",
"plugins",
]
[profile.dev]
split-debuginfo = "unpacked"

View File

@@ -1,9 +1,21 @@
[![Build Status](https://github.com/google/tarpc/workflows/Continuous%20integration/badge.svg)](https://github.com/google/tarpc/actions?query=workflow%3A%22Continuous+integration%22)
[![Latest Version](https://img.shields.io/crates/v/tarpc.svg)](https://crates.io/crates/tarpc)
[![Chat on Discord](https://img.shields.io/discord/647529123996237854)](https://discordapp.com/channels/647529123996237854)
[![Crates.io][crates-badge]][crates-url]
[![MIT licensed][mit-badge]][mit-url]
[![Build status][gh-actions-badge]][gh-actions-url]
[![Discord chat][discord-badge]][discord-url]
[crates-badge]: https://img.shields.io/crates/v/tarpc.svg
[crates-url]: https://crates.io/crates/tarpc
[mit-badge]: https://img.shields.io/badge/license-MIT-blue.svg
[mit-url]: LICENSE
[gh-actions-badge]: https://github.com/google/tarpc/workflows/Continuous%20integration/badge.svg
[gh-actions-url]: https://github.com/google/tarpc/actions?query=workflow%3A%22Continuous+integration%22
[discord-badge]: https://img.shields.io/discord/647529123996237854.svg?logo=discord&style=flat-square
[discord-url]: https://discord.gg/gXwpdSt
# tarpc
<!-- cargo-sync-readme start -->
*Disclaimer*: This is not an official Google product.
tarpc is an RPC framework for rust with a focus on ease of use. Defining a
@@ -12,7 +24,7 @@ writing a server is taken care of for you.
[Documentation](https://docs.rs/crate/tarpc/)
### What is an RPC framework?
## What is an RPC framework?
"RPC" stands for "Remote Procedure Call," a function call where the work of
producing the return value is being done somewhere else. When an rpc function is
invoked, behind the scenes the function contacts some other process somewhere
@@ -28,9 +40,9 @@ rather than in a separate language such as .proto. This means there's no separat
process, and no context switching between different languages.
Some other features of tarpc:
- Pluggable transport: any type impling `Stream<Item = Request> + Sink<Response>` can be
- Pluggable transport: any type implementing `Stream<Item = Request> + Sink<Response>` can be
used as a transport to connect the client and server.
- `Send` optional: if the transport doesn't require it, neither does tarpc!
- `Send + 'static` optional: if the transport doesn't require it, neither does tarpc!
- Cascading cancellation: dropping a request will send a cancellation message to the server.
The server will cease any unfinished work on the request, subsequently cancelling any of its
own requests, repeating for the entire chain of transitive dependencies.
@@ -39,29 +51,39 @@ Some other features of tarpc:
requests sent by the server that use the request context will propagate the request deadline.
For example, if a server is handling a request with a 10s deadline, does 2s of work, then
sends a request to another server, that server will see an 8s deadline.
- Distributed tracing: tarpc is instrumented with
[tracing](https://github.com/tokio-rs/tracing) primitives extended with
[OpenTelemetry](https://opentelemetry.io/) traces. Using a compatible tracing subscriber like
[Jaeger](https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-jaeger),
each RPC can be traced through the client, server, and other dependencies downstream of the
server. Even for applications not connected to a distributed tracing collector, the
instrumentation can also be ingested by regular loggers like
[env_logger](https://github.com/env-logger-rs/env_logger/).
- Serde serialization: enabling the `serde1` Cargo feature will make service requests and
responses `Serialize + Deserialize`. It's entirely optional, though: in-memory transports can
be used, as well, so the price of serialization doesn't have to be paid when it's not needed.
### Usage
## Usage
Add to your `Cargo.toml` dependencies:
```toml
tarpc = { version = "0.18.0", features = ["full"] }
tarpc = "0.34"
```
The `tarpc::service` attribute expands to a collection of items that form an rpc service.
These generated types make it easy and ergonomic to write servers with less boilerplate.
Simply implement the generated service trait, and you're off to the races!
### Example
## Example
For this example, in addition to tarpc, also add two other dependencies to
This example uses [tokio](https://tokio.rs), so add the following dependencies to
your `Cargo.toml`:
```toml
anyhow = "1.0"
futures = "0.3"
tokio = "0.2"
tarpc = { version = "0.31", features = ["tokio1"] }
tokio = { version = "1.0", features = ["rt-multi-thread", "macros"] }
```
In the following example, we use an in-process channel for communication between
@@ -71,15 +93,11 @@ For a more real-world example, see [example-service](example-service).
First, let's set up the dependencies and service definition.
```rust
use futures::{
future::{self, Ready},
prelude::*,
};
use futures::future::{self, Ready};
use tarpc::{
client, context,
server::{self, Handler},
server::{self, Channel},
};
use std::io;
// This is the service definition. It looks a lot like a trait definition.
// It defines one RPC, hello, which takes one arg, name, and returns a String.
@@ -100,54 +118,45 @@ implement it for our Server struct.
struct HelloServer;
impl World for HelloServer {
// Each defined rpc generates two items in the trait, a fn that serves the RPC, and
// an associated type representing the future output by the fn.
type HelloFut = Ready<String>;
fn hello(self, _: context::Context, name: String) -> Self::HelloFut {
future::ready(format!("Hello, {}!", name))
async fn hello(self, _: context::Context, name: String) -> String {
format!("Hello, {name}!")
}
}
```
Lastly let's write our `main` that will start the server. While this example uses an
[in-process
channel](https://docs.rs/tarpc/0.18.0/tarpc/transport/channel/struct.UnboundedChannel.html),
tarpc also ships bincode and JSON
tokio-net based TCP transports that are generic over all serializable types.
[in-process channel](transport::channel), tarpc also ships a generic [`serde_transport`]
behind the `serde-transport` feature, with additional [TCP](serde_transport::tcp) functionality
available behind the `tcp` feature.
```rust
#[tokio::main]
async fn main() -> io::Result<()> {
async fn main() -> anyhow::Result<()> {
let (client_transport, server_transport) = tarpc::transport::channel::unbounded();
let server = server::new(server::Config::default())
// incoming() takes a stream of transports such as would be returned by
// TcpListener::incoming (but a stream instead of an iterator).
.incoming(stream::once(future::ready(server_transport)))
.respond_with(HelloServer.serve());
let server = server::BaseChannel::with_defaults(server_transport);
tokio::spawn(server.execute(HelloServer.serve()));
tokio::spawn(server);
// WorldClient is generated by the macro. It has a constructor `new` that takes a config and
// any Transport as input
let mut client = WorldClient::new(client::Config::default(), client_transport).spawn()?;
// WorldClient is generated by the #[tarpc::service] attribute. It has a constructor `new`
// that takes a config and any Transport as input.
let client = WorldClient::new(client::Config::default(), client_transport).spawn();
// The client has an RPC method for each RPC defined in the annotated trait. It takes the same
// args as defined, with the addition of a Context, which is always the first arg. The Context
// specifies a deadline and trace information which can be helpful in debugging requests.
let hello = client.hello(context::current(), "Stim".to_string()).await?;
println!("{}", hello);
println!("{hello}");
Ok(())
}
```
### Service Documentation
## Service Documentation
Use `cargo doc` as you normally would to see the documentation created for all
items expanded by a `service!` invocation.
<!-- cargo-sync-readme end -->
License: MIT

View File

@@ -1,3 +1,388 @@
## 0.34.0 (2023-12-29)
### Breaking Changes
- `#[tarpc::server]` is no more! Service traits now use async fns.
- `Channel::execute` no longer spawns request handlers. Async-fn-in-traits makes it impossible to
add a Send bound to the future returned by `Serve::serve`. Instead, `Channel::execute` returns a
stream of futures, where each future is a request handler. To achieve the former behavior:
```rust
channel.execute(server.serve())
.for_each(|rpc| { tokio::spawn(rpc); })
```
### New Features
- Request hooks are added to the serve trait, so that it's easy to hook in cross-cutting
functionality like throttling, authorization, etc.
- The Client trait is back! This makes it possible to hook in generic client functionality like load
balancing, retries, etc.
## 0.33.0 (2023-04-01)
### Breaking Changes
Opentelemetry dependency version increased to 0.18.
## 0.32.0 (2023-03-24)
### Breaking Changes
- As part of a fix to return more channel errors in RPC results, a few error types have changed:
0. `client::RpcError::Disconnected` was split into the following errors:
- Shutdown: the client was shutdown, either intentionally or due to an error. If due to an
error, pending RPCs should see the more specific errors below.
- Send: an RPC message failed to send over the transport. Only the RPC that failed to be sent
will see this error.
- Receive: a fatal error occurred while receiving from the transport. All in-flight RPCs will
receive this error.
0. `client::ChannelError` and `server::ChannelError` are unified in `tarpc::ChannelError`.
Previously, server transport errors would not indicate during which activity the transport
error occurred. Now, just like the client already was, it will be specific: reading, readying,
sending, flushing, or closing.
## 0.31.0 (2022-11-03)
### New Features
This release adds Unix Domain Sockets to the `serde_transport` module.
To use it, enable the "unix" feature. See the docs for more information.
## 0.30.0 (2022-08-12)
### Breaking Changes
- Some types that impl Future are now annotated with `#[must_use]`. Code that previously created
these types but did not use them will now receive a warning. Code that disallows warnings will
receive a compilation error.
### Fixes
- Servers will more reliably clean up request state for requests with long deadlines when response
processing is aborted without sending a response.
### Other Changes
- `TrackedRequest` now contains a response guard that can be used to ensure state cleanup for
aborted requests. (This was already handled automatically by `InFlightRequests`).
- When the feature serde-transport is enabled, the crate tokio_serde is now re-exported.
## 0.29.0 (2022-05-26)
### Breaking Changes
`Context.deadline` is now serialized as a Duration. This prevents clock skew from affecting deadline
behavior. For more details see https://github.com/google/tarpc/pull/367 and its [related
issue](https://github.com/google/tarpc/issues/366).
## 0.28.0 (2022-04-06)
### Breaking Changes
- The minimum supported Rust version has increased to 1.58.0.
- The version of opentelemetry depended on by tarpc has increased to 0.17.0.
## 0.27.2 (2021-10-08)
### Fixes
Clients will now close their transport before dropping it. An attempt at a clean shutdown can help
the server drop its connections more quickly.
## 0.27.1 (2021-09-22)
### Breaking Changes
#### RPC error type is changing
RPC return types are changing from `Result<Response, io::Error>` to `Result<Response,
tarpc::client::RpcError>`.
Becaue tarpc is a library, not an application, it should strive to
use structured errors in its API so that users have maximal flexibility
in how they handle errors. io::Error makes that hard, because it is a
kitchen-sink error type.
RPCs in particular only have 3 classes of errors:
- The connection breaks.
- The request expires.
- The server decides not to process the request.
RPC responses can also contain application-specific errors, but from the
perspective of the RPC library, those are opaque to the framework, classified
as successful responsees.
### Open Telemetry
The Opentelemetry dependency is updated to version 0.16.x.
## 0.27.0 (2021-09-22)
This version was yanked due to tarpc-plugins version mismatches.
## 0.26.0 (2021-04-14)
### New Features
#### Tracing
tarpc is now instrumented with tracing primitives extended with
OpenTelemetry traces. Using a compatible tracing-opentelemetry
subscriber like Jaeger, each RPC can be traced through the client,
server, amd other dependencies downstream of the server. Even for
applications not connected to a distributed tracing collector, the
instrumentation can also be ingested by regular loggers like env_logger.
### Breaking Changes
#### Logging
Logged events are now structured using tracing. For applications using a
logger and not a tracing subscriber, these logs may look different or
contain information in a less consumable manner. The easiest solution is
to add a tracing subscriber that logs to stdout, such as
tracing_subscriber::fmt.
#### Context
- Context no longer has parent_span, which was actually never needed,
because the context sent in an RPC is inherently the parent context.
For purposes of distributed tracing, the client side of the RPC has all
necessary information to link the span to its parent; the server side
need do nothing more than export the (trace ID, span ID) tuple.
- Context has a new field, SamplingDecision, which has two variants,
Sampled and Unsampled. This field can be used by downstream systems to
determine whether a trace needs to be exported. If the parent span is
sampled, the expectation is that all child spans be exported, as well;
to do otherwise could result in lossy traces being exported. Note that
if an Openetelemetry tracing subscriber is not installed, the fallback
context will still be used, but the Context's sampling decision will
always be inherited by the parent Context's sampling decision.
- Context::scope has been removed. Context propagation is now done via
tracing's task-local spans. Spans can be propagated across tasks via
Span::in_scope. When a service receives a request, it attaches an
Opentelemetry context to the local Span created before request handling,
and this context contains the request deadline. This span-local deadline
is retrieved by Context::current, but it cannot be modified so that
future Context::current calls contain a different deadline. However, the
deadline in the context passed into an RPC call will override it, so
users can retrieve the current context and then modify the deadline
field, as has been historically possible.
- Context propgation precedence changes: when an RPC is initiated, the
current Span's Opentelemetry context takes precedence over the trace
context passed into the RPC method. If there is no current Span, then
the trace context argument is used as it has been historically. Note
that Opentelemetry context propagation requires an Opentelemetry
tracing subscriber to be installed.
#### Server
- The server::Channel trait now has an additional required associated
type and method which returns the underlying transport. This makes it
more ergonomic for users to retrieve transport-specific information,
like IP Address. BaseChannel implements Channel::transport by returning
the underlying transport, and channel decorators like Throttler just
delegate to the Channel::transport method of the wrapped channel.
#### Client
- NewClient::spawn no longer returns a result, as spawn can't fail.
### References
1. https://github.com/tokio-rs/tracing
2. https://opentelemetry.io
3. https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-jaeger
4. https://github.com/env-logger-rs/env_logger
## 0.25.0 (2021-03-10)
### Breaking Changes
#### Major server module refactoring
1. Renames
Some of the items in this module were renamed to be less generic:
- Handler => Incoming
- ClientHandler => Requests
- ResponseHandler => InFlightRequest
- Channel::{respond_with => requests}
In the case of Handler: handler of *what*? Now it's a bit clearer that this is a stream of Channels
(aka *incoming* connections).
Similarly, ClientHandler was a stream of requests over a single connection. Hopefully Requests
better reflects that.
ResponseHandler was renamed InFlightRequest because it no longer contains the serving function.
Instead, it is just the request, plus the response channel and an abort hook. As a result of this,
Channel::respond_with underwent a big change: it used to take the serving function and return a
ClientHandler; now it has been renamed Channel::requests and does not take any args.
2. Execute methods
All methods thats actually result in responses being generated have been consolidated into methods
named `execute`:
- InFlightRequest::execute returns a future that completes when a response has been generated and
sent to the server Channel.
- Requests::execute automatically spawns response handlers for all requests over a single channel.
- Channel::execute is a convenience for `channel.requests().execute()`.
- Incoming::execute automatically spawns response handlers for all requests over all channels.
3. Removal of Server.
server::Server was removed, as it provided no value over the Incoming/Channel abstractions.
Additionally, server::new was removed, since it just returned a Server.
#### Client RPC methods now take &self
This required the breaking change of removing the Client trait. The intent of the Client trait was
to facilitate the decorator pattern by allowing users to create their own Clients that added
behavior on top of the base client. Unfortunately, this trait had become a maintenance burden,
consistently causing issues with lifetimes and the lack of generic associated types. Specifically,
it meant that Client impls could not use async fns, which is no longer tenable today, with channel
libraries moving to async fns.
#### Servers no longer send deadline-exceed responses.
The deadline-exceeded response was largely redundant, because the client
shouldn't normally be waiting for such a response, anyway -- the normal
client will automatically remove the in-flight request when it reaches
the deadline.
This also allows for internalizing the expiration+cleanup logic entirely
within BaseChannel, without having it leak into the Channel trait and
requiring action taken by the Requests struct.
#### Clients no longer send cancel messages when the request deadline is exceeded.
The server already knows when the request deadline was exceeded, so the client didn't need to inform
it.
### Fixes
- When a channel is dropped, all in-flight requests for that channel are now aborted.
## 0.24.1 (2020-12-28)
### Breaking Changes
Upgrades tokio to 1.0.
## 0.24.0 (2020-12-28)
This release was yanked.
## 0.23.0 (2020-10-19)
### Breaking Changes
Upgrades tokio to 0.3.
## 0.22.0 (2020-08-02)
This release adds some flexibility and consistency to `serde_transport`, with one new feature and
one small breaking change.
### New Features
`serde_transport::tcp` now exposes framing configuration on `connect()` and `listen()`. This is
useful if, for instance, you want to send requests or responses that are larger than the maximum
payload allowed by default:
```rust
let mut transport = tarpc::serde_transport::tcp::connect(server_addr, Json::default);
transport.config_mut().max_frame_length(4294967296);
let mut client = MyClient::new(client::Config::default(), transport.await?).spawn()?;
```
### Breaking Changes
The codec argument to `serde_transport::tcp::connect` changed from a Codec to impl Fn() -> Codec,
to be consistent with `serde_transport::tcp::listen`. While only one Codec is needed, more than one
person has been tripped up by the inconsistency between `connect` and `listen`. Unfortunately, the
compiler errors are not much help in this case, so it was decided to simply do the more intuitive
thing so that the compiler doesn't need to step in in the first place.
## 0.21.1 (2020-08-02)
### New Features
#### #[tarpc::server] diagnostics
When a service impl uses #[tarpc::server], only `async fn`s are re-written. This can lead to
confusing compiler errors about missing associated types:
```
error: not all trait items implemented, missing: `HelloFut`
--> $DIR/tarpc_server_missing_async.rs:9:1
|
9 | impl World for HelloServer {
| ^^^^
```
The proc macro now provides better diagnostics for this case:
```
error: not all trait items implemented, missing: `HelloFut`
--> $DIR/tarpc_server_missing_async.rs:9:1
|
9 | impl World for HelloServer {
| ^^^^
error: hint: `#[tarpc::server]` only rewrites async fns, and `fn hello` is not async
--> $DIR/tarpc_server_missing_async.rs:10:5
|
10 | fn hello(name: String) -> String {
| ^^
```
### Bug Fixes
#### Fixed client hanging when server shuts down
Previously, clients would ignore when the read half of the transport was closed, continuing to
write requests. This didn't make much sense, because without the ability to receive responses,
clients have no way to know if requests were actually processed by the server. It basically just
led to clients that would hang for a few seconds before shutting down. This has now been
corrected: clients will immediately shut down when the read-half of the transport is closed.
#### More docs.rs documentation
Previously, docs.rs only documented items enabled by default, notably leaving out documentation
for tokio and serde features. This has now been corrected: docs.rs should have documentation
for all optional features.
## 0.21.0 (2020-06-26)
### New Features
A new proc macro, `#[tarpc::server]` was added! This enables service impls to elide the boilerplate
of specifying associated types for each RPC. With the ubiquity of async-await, most code won't have
nameable futures and will just be boxing the return type anyway. This macro does that for you.
### Breaking Changes
- Enums had `_non_exhaustive` fields replaced with the #[non_exhaustive] attribute.
### Bug Fixes
- https://github.com/google/tarpc/issues/304
A race condition in code that limits number of connections per client caused occasional panics.
- https://github.com/google/tarpc/pull/295
Made request timeouts account for time spent in the outbound buffer. Previously, a large outbound
queue would lead to requests not timing out correctly.
## 0.20.0 (2019-12-11)
### Breaking Changes
@@ -10,7 +395,7 @@
## 0.13.0 (2018-10-16)
### Breaking Changes
### Breaking Changes
Version 0.13 marks a significant departure from previous versions of tarpc. The
API has changed significantly. The tokio-proto crate has been torn out and

View File

@@ -1,8 +1,9 @@
[package]
name = "tarpc-example-service"
version = "0.6.0"
version = "0.15.0"
rust-version = "1.56"
authors = ["Tim Kuehn <tikue@google.com>"]
edition = "2018"
edition = "2021"
license = "MIT"
documentation = "https://docs.rs/tarpc-example-service"
homepage = "https://github.com/google/tarpc"
@@ -13,13 +14,19 @@ readme = "../README.md"
description = "An example server built on tarpc."
[dependencies]
clap = "2.0"
anyhow = "1.0"
clap = { version = "3.0.0-rc.9", features = ["derive"] }
log = "0.4"
futures = "0.3"
serde = { version = "1.0" }
tarpc = { version = "0.20", path = "../tarpc", features = ["full"] }
tokio = { version = "0.2", features = ["full"] }
tokio-serde = { version = "0.6", features = ["json"] }
env_logger = "0.6"
opentelemetry = { version = "0.21.0" }
opentelemetry-jaeger = { version = "0.20.0", features = ["rt-tokio"] }
rand = "0.8"
tarpc = { version = "0.34", path = "../tarpc", features = ["full"] }
tokio = { version = "1", features = ["macros", "net", "rt-multi-thread"] }
tracing = { version = "0.1" }
tracing-opentelemetry = "0.22.0"
tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
opentelemetry_sdk = "0.21.1"
[lib]
name = "service"

15
example-service/README.md Normal file
View File

@@ -0,0 +1,15 @@
# Example
Example service to demonstrate how to set up `tarpc` with [Jaeger](https://www.jaegertracing.io). To see traces Jaeger, run the following with `RUST_LOG=trace`.
## Server
```bash
cargo run --bin server -- --port 50051
```
## Client
```bash
cargo run --bin client -- --server-addr "[::1]:50051" --name "Bob"
```

View File

@@ -4,55 +4,53 @@
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use clap::{App, Arg};
use std::{io, net::SocketAddr};
use tarpc::{client, context};
use tokio_serde::formats::Json;
use clap::Parser;
use service::{init_tracing, WorldClient};
use std::{net::SocketAddr, time::Duration};
use tarpc::{client, context, tokio_serde::formats::Json};
use tokio::time::sleep;
use tracing::Instrument;
#[derive(Parser)]
struct Flags {
/// Sets the server address to connect to.
#[clap(long)]
server_addr: SocketAddr,
/// Sets the name to say hello to.
#[clap(long)]
name: String,
}
#[tokio::main]
async fn main() -> io::Result<()> {
let flags = App::new("Hello Client")
.version("0.1")
.author("Tim <tikue@google.com>")
.about("Say hello!")
.arg(
Arg::with_name("server_addr")
.long("server_addr")
.value_name("ADDRESS")
.help("Sets the server address to connect to.")
.required(true)
.takes_value(true),
)
.arg(
Arg::with_name("name")
.short("n")
.long("name")
.value_name("STRING")
.help("Sets the name to say hello to.")
.required(true)
.takes_value(true),
)
.get_matches();
async fn main() -> anyhow::Result<()> {
let flags = Flags::parse();
init_tracing("Tarpc Example Client")?;
let server_addr = flags.value_of("server_addr").unwrap();
let server_addr = server_addr
.parse::<SocketAddr>()
.unwrap_or_else(|e| panic!(r#"--server_addr value "{}" invalid: {}"#, server_addr, e));
let name = flags.value_of("name").unwrap().into();
let transport = tarpc::serde_transport::tcp::connect(server_addr, Json::default()).await?;
let mut transport = tarpc::serde_transport::tcp::connect(flags.server_addr, Json::default);
transport.config_mut().max_frame_length(usize::MAX);
// WorldClient is generated by the service attribute. It has a constructor `new` that takes a
// config and any Transport as input.
let mut client = service::WorldClient::new(client::Config::default(), transport).spawn()?;
let client = WorldClient::new(client::Config::default(), transport.await?).spawn();
// The client has an RPC method for each RPC defined in the annotated trait. It takes the same
// args as defined, with the addition of a Context, which is always the first arg. The Context
// specifies a deadline and trace information which can be helpful in debugging requests.
let hello = client.hello(context::current(), name).await?;
let hello = async move {
// Send the request twice, just to be safe! ;)
tokio::select! {
hello1 = client.hello(context::current(), format!("{}1", flags.name)) => { hello1 }
hello2 = client.hello(context::current(), format!("{}2", flags.name)) => { hello2 }
}
}
.instrument(tracing::info_span!("Two Hellos"))
.await;
println!("{}", hello);
match hello {
Ok(hello) => tracing::info!("{hello:?}"),
Err(e) => tracing::warn!("{:?}", anyhow::Error::from(e)),
}
// Let the background span processor finish.
sleep(Duration::from_micros(1)).await;
opentelemetry::global::shutdown_tracer_provider();
Ok(())
}

View File

@@ -4,6 +4,9 @@
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use std::env;
use tracing_subscriber::{fmt::format::FmtSpan, prelude::*};
/// This is the service definition. It looks a lot like a trait definition.
/// It defines one RPC, hello, which takes one arg, name, and returns a String.
#[tarpc::service]
@@ -11,3 +14,21 @@ pub trait World {
/// Returns a greeting for name.
async fn hello(name: String) -> String;
}
/// Initializes an OpenTelemetry tracing subscriber with a Jaeger backend.
pub fn init_tracing(service_name: &str) -> anyhow::Result<()> {
env::set_var("OTEL_BSP_MAX_EXPORT_BATCH_SIZE", "12");
let tracer = opentelemetry_jaeger::new_agent_pipeline()
.with_service_name(service_name)
.with_max_packet_size(2usize.pow(13))
.install_batch(opentelemetry_sdk::runtime::Tokio)?;
tracing_subscriber::registry()
.with(tracing_subscriber::EnvFilter::from_default_env())
.with(tracing_subscriber::fmt::layer().with_span_events(FmtSpan::NEW | FmtSpan::CLOSE))
.with(tracing_opentelemetry::layer().with_tracer(tracer))
.try_init()?;
Ok(())
}

View File

@@ -4,21 +4,30 @@
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use clap::{App, Arg};
use futures::{
future::{self, Ready},
prelude::*,
use clap::Parser;
use futures::{future, prelude::*};
use rand::{
distributions::{Distribution, Uniform},
thread_rng,
};
use service::World;
use service::{init_tracing, World};
use std::{
io,
net::{IpAddr, SocketAddr},
net::{IpAddr, Ipv6Addr, SocketAddr},
time::Duration,
};
use tarpc::{
context,
server::{self, Channel, Handler},
server::{self, incoming::Incoming, Channel},
tokio_serde::formats::Json,
};
use tokio_serde::formats::Json;
use tokio::time;
#[derive(Parser)]
struct Flags {
/// Sets the port number to listen on.
#[clap(long)]
port: u16,
}
// This is the type that implements the generated World trait. It is the business logic
// and is used to start the server.
@@ -26,59 +35,41 @@ use tokio_serde::formats::Json;
struct HelloServer(SocketAddr);
impl World for HelloServer {
// Each defined rpc generates two items in the trait, a fn that serves the RPC, and
// an associated type representing the future output by the fn.
type HelloFut = Ready<String>;
fn hello(self, _: context::Context, name: String) -> Self::HelloFut {
future::ready(format!(
"Hello, {}! You are connected from {:?}.",
name, self.0
))
async fn hello(self, _: context::Context, name: String) -> String {
let sleep_time =
Duration::from_millis(Uniform::new_inclusive(1, 10).sample(&mut thread_rng()));
time::sleep(sleep_time).await;
format!("Hello, {name}! You are connected from {}", self.0)
}
}
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
#[tokio::main]
async fn main() -> io::Result<()> {
env_logger::init();
async fn main() -> anyhow::Result<()> {
let flags = Flags::parse();
init_tracing("Tarpc Example Server")?;
let flags = App::new("Hello Server")
.version("0.1")
.author("Tim <tikue@google.com>")
.about("Say hello!")
.arg(
Arg::with_name("port")
.short("p")
.long("port")
.value_name("NUMBER")
.help("Sets the port number to listen on")
.required(true)
.takes_value(true),
)
.get_matches();
let port = flags.value_of("port").unwrap();
let port = port
.parse()
.unwrap_or_else(|e| panic!(r#"--port value "{}" invalid: {}"#, port, e));
let server_addr = (IpAddr::from([0, 0, 0, 0]), port);
let server_addr = (IpAddr::V6(Ipv6Addr::LOCALHOST), flags.port);
// JSON transport is provided by the json_transport tarpc module. It makes it easy
// to start up a serde-powered json serialization strategy over TCP.
tarpc::serde_transport::tcp::listen(&server_addr, Json::default)
.await?
let mut listener = tarpc::serde_transport::tcp::listen(&server_addr, Json::default).await?;
tracing::info!("Listening on port {}", listener.local_addr().port());
listener.config_mut().max_frame_length(usize::MAX);
listener
// Ignore accept errors.
.filter_map(|r| future::ready(r.ok()))
.map(server::BaseChannel::with_defaults)
// Limit channels to 1 per IP.
.max_channels_per_key(1, |t| t.as_ref().peer_addr().unwrap().ip())
.max_channels_per_key(1, |t| t.transport().peer_addr().unwrap().ip())
// serve is generated by the service attribute. It takes as input any type implementing
// the generated World trait.
.map(|channel| {
let server = HelloServer(channel.as_ref().as_ref().peer_addr().unwrap());
channel.respond_with(server.serve()).execute()
let server = HelloServer(channel.transport().peer_addr().unwrap());
channel.execute(server.serve()).for_each(spawn)
})
// Max 10 channels.
.buffer_unordered(10)

View File

@@ -67,7 +67,7 @@ else
fi
printf "${PREFIX} Checking for rustfmt ... "
command -v cargo fmt &>/dev/null
command -v rustfmt &>/dev/null
if [ $? == 0 ]; then
printf "${SUCCESS}\n"
else
@@ -93,19 +93,19 @@ diff=""
for file in $(git diff --name-only --cached);
do
if [ ${file: -3} == ".rs" ]; then
diff="$diff$(cargo fmt -- --skip-children --write-mode=diff $file)"
diff="$diff$(rustfmt --edition 2018 --check $file)"
if [ $? != 0 ]; then
FMTRESULT=1
fi
fi
done
if grep --quiet "^[-+]" <<< "$diff"; then
FMTRESULT=1
fi
if [ "${TARPC_SKIP_RUSTFMT}" == 1 ]; then
printf "${SKIPPED}\n"$?
elif [ ${FMTRESULT} != 0 ]; then
FAILED=1
printf "${FAILURE}\n"
echo "$diff" | sed 's/Using rustfmt config file.*$/d/'
echo "$diff"
else
printf "${SUCCESS}\n"
fi

View File

@@ -84,18 +84,19 @@ command -v rustup &>/dev/null
if [ "$?" == 0 ]; then
printf "${SUCCESS}\n"
try_run "Building ... " cargo +stable build --color=always
try_run "Testing ... " cargo +stable test --color=always
try_run "Testing with all features enabled ... " cargo +stable test --all-features --color=always
for EXAMPLE in $(cargo +stable run --example 2>&1 | grep ' ' | awk '{print $1}')
do
try_run "Running example \"$EXAMPLE\" ... " cargo +stable run --example $EXAMPLE
done
check_toolchain nightly
if [ ${TOOLCHAIN_RESULT} == 1 ]; then
exit 1
if [ ${TOOLCHAIN_RESULT} != 1 ]; then
try_run "Running clippy ... " cargo +nightly clippy --color=always -Z unstable-options -- --deny warnings
fi
try_run "Building ... " cargo build --color=always
try_run "Testing ... " cargo test --color=always
try_run "Testing with all features enabled ... " cargo test --all-features --color=always
for EXAMPLE in $(cargo run --example 2>&1 | grep ' ' | awk '{print $1}')
do
try_run "Running example \"$EXAMPLE\" ... " cargo run --example $EXAMPLE
done
fi

View File

@@ -1,8 +1,9 @@
[package]
name = "tarpc-plugins"
version = "0.7.0"
version = "0.13.0"
rust-version = "1.56"
authors = ["Adam Wright <adam.austin.wright@gmail.com>", "Tim Kuehn <timothy.j.kuehn@gmail.com>"]
edition = "2018"
edition = "2021"
license = "MIT"
documentation = "https://docs.rs/tarpc-plugins"
homepage = "https://github.com/google/tarpc"
@@ -19,14 +20,15 @@ serde1 = []
travis-ci = { repository = "google/tarpc" }
[dependencies]
syn = { version = "1.0.11", features = ["full"] }
quote = "1.0.2"
proc-macro2 = "1.0.6"
proc-macro2 = "1.0"
quote = "1.0"
syn = { version = "1.0", features = ["full"] }
[lib]
proc-macro = true
[dev-dependencies]
assert-type-eq = "0.1.0"
futures = "0.3"
serde = { version = "1.0", features = ["derive"] }
tarpc = { path = "../tarpc" }
tarpc = { path = "../tarpc", features = ["serde1"] }

9
plugins/LICENSE Normal file
View File

@@ -0,0 +1,9 @@
The MIT License (MIT)
Copyright 2016 Google Inc. All Rights Reserved.
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

View File

@@ -12,17 +12,32 @@ extern crate quote;
extern crate syn;
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote, ToTokens};
use syn::{
braced, parenthesized,
braced,
ext::IdentExt,
parenthesized,
parse::{Parse, ParseStream},
parse_macro_input, parse_quote,
punctuated::Punctuated,
spanned::Spanned,
token::Comma,
Attribute, FnArg, Ident, Lit, LitBool, MetaNameValue, Pat, PatType, ReturnType, Token, Type,
Visibility,
};
/// Accumulates multiple errors into a result.
/// Only use this for recoverable errors, i.e. non-parse errors. Fatal errors should early exit to
/// avoid further complications.
macro_rules! extend_errors {
($errors: ident, $e: expr) => {
match $errors {
Ok(_) => $errors = Err($e),
Err(ref mut errors) => errors.extend($e),
}
};
}
struct Service {
attrs: Vec<Attribute>,
vis: Visibility,
@@ -42,27 +57,39 @@ impl Parse for Service {
let attrs = input.call(Attribute::parse_outer)?;
let vis = input.parse()?;
input.parse::<Token![trait]>()?;
let ident = input.parse()?;
let ident: Ident = input.parse()?;
let content;
braced!(content in input);
let mut rpcs = Vec::<RpcMethod>::new();
while !content.is_empty() {
rpcs.push(content.parse()?);
}
let mut ident_errors = Ok(());
for rpc in &rpcs {
if rpc.ident == "new" {
return Err(input.error(format!(
"method name conflicts with generated fn `{}Client::new`",
ident
)));
extend_errors!(
ident_errors,
syn::Error::new(
rpc.ident.span(),
format!(
"method name conflicts with generated fn `{}Client::new`",
ident.unraw()
)
)
);
}
if rpc.ident == "serve" {
return Err(input.error(format!(
"method name conflicts with generated fn `{}::serve`",
ident
)));
extend_errors!(
ident_errors,
syn::Error::new(
rpc.ident.span(),
format!("method name conflicts with generated fn `{ident}::serve`")
)
);
}
}
ident_errors?;
Ok(Self {
attrs,
vis,
@@ -80,17 +107,28 @@ impl Parse for RpcMethod {
let ident = input.parse()?;
let content;
parenthesized!(content in input);
let args: Punctuated<FnArg, Comma> = content.parse_terminated(FnArg::parse)?;
let args = args
.into_iter()
.map(|arg| match arg {
FnArg::Typed(captured) => match *captured.pat {
Pat::Ident(_) => Ok(captured),
_ => Err(input.error("patterns aren't allowed in RPC args")),
},
FnArg::Receiver(_) => Err(input.error("method args cannot start with self")),
})
.collect::<Result<_, _>>()?;
let mut args = Vec::new();
let mut errors = Ok(());
for arg in content.parse_terminated::<FnArg, Comma>(FnArg::parse)? {
match arg {
FnArg::Typed(captured) if matches!(&*captured.pat, Pat::Ident(_)) => {
args.push(captured);
}
FnArg::Typed(captured) => {
extend_errors!(
errors,
syn::Error::new(captured.pat.span(), "patterns aren't allowed in RPC args")
);
}
FnArg::Receiver(_) => {
extend_errors!(
errors,
syn::Error::new(arg.span(), "method args cannot start with self")
);
}
}
}
errors?;
let output = input.parse()?;
input.parse::<Token![;]>()?;
@@ -109,35 +147,93 @@ struct DeriveSerde(bool);
impl Parse for DeriveSerde {
fn parse(input: ParseStream) -> syn::Result<Self> {
if input.is_empty() {
return Ok(Self(cfg!(feature = "serde1")));
}
match input.parse::<MetaNameValue>()? {
MetaNameValue {
ref path, ref lit, ..
} if path.segments.len() == 1
&& path.segments.first().unwrap().ident == "derive_serde" =>
{
match lit {
Lit::Bool(LitBool { value: true, .. }) if cfg!(feature = "serde1") => {
Ok(Self(true))
}
Lit::Bool(LitBool { value: true, .. }) => {
Err(input
.error("To enable serde, first enable the `serde1` feature of tarpc"))
}
Lit::Bool(LitBool { value: false, .. }) => Ok(Self(false)),
_ => Err(input.error("`derive_serde` expects a value of type `bool`")),
let mut result = Ok(None);
let mut derive_serde = Vec::new();
let meta_items = input.parse_terminated::<MetaNameValue, Comma>(MetaNameValue::parse)?;
for meta in meta_items {
if meta.path.segments.len() != 1 {
extend_errors!(
result,
syn::Error::new(
meta.span(),
"tarpc::service does not support this meta item"
)
);
continue;
}
let segment = meta.path.segments.first().unwrap();
if segment.ident != "derive_serde" {
extend_errors!(
result,
syn::Error::new(
meta.span(),
"tarpc::service does not support this meta item"
)
);
continue;
}
match meta.lit {
Lit::Bool(LitBool { value: true, .. }) if cfg!(feature = "serde1") => {
result = result.and(Ok(Some(true)))
}
Lit::Bool(LitBool { value: true, .. }) => {
extend_errors!(
result,
syn::Error::new(
meta.span(),
"To enable serde, first enable the `serde1` feature of tarpc"
)
);
}
Lit::Bool(LitBool { value: false, .. }) => result = result.and(Ok(Some(false))),
_ => extend_errors!(
result,
syn::Error::new(
meta.lit.span(),
"`derive_serde` expects a value of type `bool`"
)
),
}
_ => {
Err(input
.error("tarpc::service only supports one meta item, `derive_serde = {bool}`"))
derive_serde.push(meta);
}
if derive_serde.len() > 1 {
for (i, derive_serde) in derive_serde.iter().enumerate() {
extend_errors!(
result,
syn::Error::new(
derive_serde.span(),
format!(
"`derive_serde` appears more than once (occurrence #{})",
i + 1
)
)
);
}
}
let derive_serde = result?.unwrap_or(cfg!(feature = "serde1"));
Ok(Self(derive_serde))
}
}
/// A helper attribute to avoid a direct dependency on Serde.
///
/// Adds the following annotations to the annotated item:
///
/// ```rust
/// #[derive(::tarpc::serde::Serialize, ::tarpc::serde::Deserialize)]
/// #[serde(crate = "tarpc::serde")]
/// # struct Foo;
/// ```
#[proc_macro_attribute]
pub fn derive_serde(_attr: TokenStream, item: TokenStream) -> TokenStream {
let mut gen: proc_macro2::TokenStream = quote! {
#[derive(::tarpc::serde::Serialize, ::tarpc::serde::Deserialize)]
#[serde(crate = "::tarpc::serde")]
};
gen.extend(proc_macro2::TokenStream::from(item));
proc_macro::TokenStream::from(gen)
}
/// Generates:
/// - service trait
/// - serve fn
@@ -156,28 +252,40 @@ pub fn service(attr: TokenStream, input: TokenStream) -> TokenStream {
ref rpcs,
} = parse_macro_input!(input as Service);
let camel_case_fn_names: &[String] = &rpcs
let camel_case_fn_names: &Vec<_> = &rpcs
.iter()
.map(|rpc| snake_to_camel(&rpc.ident.to_string()))
.collect::<Vec<_>>();
.map(|rpc| snake_to_camel(&rpc.ident.unraw().to_string()))
.collect();
let args: &[&[PatType]] = &rpcs.iter().map(|rpc| &*rpc.args).collect::<Vec<_>>();
let response_fut_name = &format!("{}ResponseFut", ident);
let derive_serialize = quote!(#[derive(serde::Serialize, serde::Deserialize)]);
let derive_serialize = if derive_serde.0 {
Some(
quote! {#[derive(::tarpc::serde::Serialize, ::tarpc::serde::Deserialize)]
#[serde(crate = "::tarpc::serde")]},
)
} else {
None
};
let gen_args = &GenArgs {
attrs,
vis,
rpcs,
args,
response_fut_name,
let methods = rpcs.iter().map(|rpc| &rpc.ident).collect::<Vec<_>>();
let request_names = methods
.iter()
.map(|m| format!("{ident}.{m}"))
.collect::<Vec<_>>();
ServiceGenerator {
service_ident: ident,
client_stub_ident: &format_ident!("{}Stub", ident),
server_ident: &format_ident!("Serve{}", ident),
response_fut_ident: &Ident::new(&response_fut_name, ident.span()),
client_ident: &format_ident!("{}Client", ident),
request_ident: &format_ident!("{}Request", ident),
response_ident: &format_ident!("{}Response", ident),
vis,
args,
method_attrs: &rpcs.iter().map(|rpc| &*rpc.attrs).collect::<Vec<_>>(),
method_names: &rpcs.iter().map(|rpc| &rpc.ident).collect::<Vec<_>>(),
method_idents: &methods,
request_names: &request_names,
attrs,
rpcs,
return_types: &rpcs
.iter()
.map(|rpc| match rpc.output {
@@ -185,7 +293,7 @@ pub fn service(attr: TokenStream, input: TokenStream) -> TokenStream {
ReturnType::Default => unit_type,
})
.collect::<Vec<_>>(),
arg_vars: &args
arg_pats: &args
.iter()
.map(|args| args.iter().map(|arg| &*arg.pat).collect())
.collect::<Vec<_>>(),
@@ -194,375 +302,318 @@ pub fn service(attr: TokenStream, input: TokenStream) -> TokenStream {
.zip(camel_case_fn_names.iter())
.map(|(rpc, name)| Ident::new(name, rpc.ident.span()))
.collect::<Vec<_>>(),
future_idents: &camel_case_fn_names
.iter()
.map(|name| format_ident!("{}Fut", name))
.collect::<Vec<_>>(),
derive_serialize: if derive_serde.0 {
Some(&derive_serialize)
} else {
None
},
};
let tokens = vec![
trait_service(gen_args),
struct_server(gen_args),
impl_serve_for_server(gen_args),
enum_request(gen_args),
enum_response(gen_args),
enum_response_future(gen_args),
impl_debug_for_response_future(gen_args),
impl_future_for_response_future(gen_args),
struct_client(gen_args),
impl_from_for_client(gen_args),
impl_client_new(gen_args),
impl_client_rpc_methods(gen_args),
];
tokens
.into_iter()
.collect::<proc_macro2::TokenStream>()
.into()
derive_serialize: derive_serialize.as_ref(),
}
.into_token_stream()
.into()
}
// Things needed to generate the service items: trait, serve impl, request/response enums, and
// the client stub.
struct GenArgs<'a> {
attrs: &'a [Attribute],
rpcs: &'a [RpcMethod],
struct ServiceGenerator<'a> {
service_ident: &'a Ident,
client_stub_ident: &'a Ident,
server_ident: &'a Ident,
response_fut_ident: &'a Ident,
response_fut_name: &'a str,
client_ident: &'a Ident,
request_ident: &'a Ident,
response_ident: &'a Ident,
method_attrs: &'a [&'a [Attribute]],
vis: &'a Visibility,
method_names: &'a [&'a Ident],
attrs: &'a [Attribute],
rpcs: &'a [RpcMethod],
camel_case_idents: &'a [Ident],
method_idents: &'a [&'a Ident],
request_names: &'a [String],
method_attrs: &'a [&'a [Attribute]],
args: &'a [&'a [PatType]],
return_types: &'a [&'a Type],
arg_vars: &'a [Vec<&'a Pat>],
camel_case_idents: &'a [Ident],
future_idents: &'a [Ident],
derive_serialize: Option<&'a proc_macro2::TokenStream>,
arg_pats: &'a [Vec<&'a Pat>],
derive_serialize: Option<&'a TokenStream2>,
}
fn trait_service(
&GenArgs {
attrs,
rpcs,
vis,
future_idents,
return_types,
service_ident,
server_ident,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
let types_and_fns = rpcs
.iter()
.zip(future_idents.iter())
.zip(return_types.iter())
.map(
|(
(
impl<'a> ServiceGenerator<'a> {
fn trait_service(&self) -> TokenStream2 {
let &Self {
attrs,
rpcs,
vis,
return_types,
service_ident,
client_stub_ident,
request_ident,
response_ident,
server_ident,
..
} = self;
let rpc_fns = rpcs
.iter()
.zip(return_types.iter())
.map(
|(
RpcMethod {
attrs, ident, args, ..
},
future_type,
),
output,
)| {
let ty_doc = format!("The response future returned by {}.", ident);
quote! {
#[doc = #ty_doc]
type #future_type: std::future::Future<Output = #output>;
output,
)| {
quote! {
#( #attrs )*
async fn #ident(self, context: ::tarpc::context::Context, #( #args ),*) -> #output;
}
},
);
#( #attrs )*
fn #ident(self, context: tarpc::context::Context, #( #args ),*) -> Self::#future_type;
}
},
);
let stub_doc = format!("The stub trait for service [`{service_ident}`].");
quote! {
#( #attrs )*
#vis trait #service_ident: ::core::marker::Sized {
#( #rpc_fns )*
quote! {
#( #attrs )*
#vis trait #service_ident: Clone {
#( #types_and_fns )*
/// Returns a serving function to use with tarpc::server::Server.
fn serve(self) -> #server_ident<Self> {
#server_ident { service: self }
}
}
}
}
fn struct_server(
&GenArgs {
vis, server_ident, ..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
#[derive(Clone)]
#vis struct #server_ident<S> {
service: S,
}
}
}
fn impl_serve_for_server(
&GenArgs {
request_ident,
server_ident,
service_ident,
response_ident,
response_fut_ident,
camel_case_idents,
arg_vars,
method_names,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
impl<S> tarpc::server::Serve<#request_ident> for #server_ident<S>
where S: #service_ident
{
type Resp = #response_ident;
type Fut = #response_fut_ident<S>;
fn serve(self, ctx: tarpc::context::Context, req: #request_ident) -> Self::Fut {
match req {
#(
#request_ident::#camel_case_idents{ #( #arg_vars ),* } => {
#response_fut_ident::#camel_case_idents(
#service_ident::#method_names(
self.service, ctx, #( #arg_vars ),*))
}
)*
/// Returns a serving function to use with
/// [InFlightRequest::execute](::tarpc::server::InFlightRequest::execute).
fn serve(self) -> #server_ident<Self> {
#server_ident { service: self }
}
}
}
}
}
fn enum_request(
&GenArgs {
derive_serialize,
vis,
request_ident,
camel_case_idents,
args,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
/// The request sent over the wire from the client to the server.
#[derive(Debug)]
#derive_serialize
#vis enum #request_ident {
#( #camel_case_idents{ #( #args ),* } ),*
}
}
}
fn enum_response(
&GenArgs {
derive_serialize,
vis,
response_ident,
camel_case_idents,
return_types,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
/// The response sent over the wire from the server to the client.
#[derive(Debug)]
#derive_serialize
#vis enum #response_ident {
#( #camel_case_idents(#return_types) ),*
}
}
}
fn enum_response_future(
&GenArgs {
vis,
service_ident,
response_fut_ident,
camel_case_idents,
future_idents,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
/// A future resolving to a server response.
#vis enum #response_fut_ident<S: #service_ident> {
#( #camel_case_idents(<S as #service_ident>::#future_idents) ),*
}
}
}
fn impl_debug_for_response_future(
&GenArgs {
service_ident,
response_fut_ident,
response_fut_name,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
impl<S: #service_ident> std::fmt::Debug for #response_fut_ident<S> {
fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
fmt.debug_struct(#response_fut_name).finish()
#[doc = #stub_doc]
#vis trait #client_stub_ident: ::tarpc::client::stub::Stub<Req = #request_ident, Resp = #response_ident> {
}
}
}
}
fn impl_future_for_response_future(
&GenArgs {
service_ident,
response_fut_ident,
response_ident,
camel_case_idents,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
impl<S: #service_ident> std::future::Future for #response_fut_ident<S> {
type Output = #response_ident;
fn poll(self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>)
-> std::task::Poll<#response_ident>
impl<S> #client_stub_ident for S
where S: ::tarpc::client::stub::Stub<Req = #request_ident, Resp = #response_ident>
{
unsafe {
match std::pin::Pin::get_unchecked_mut(self) {
}
}
}
fn struct_server(&self) -> TokenStream2 {
let &Self {
vis, server_ident, ..
} = self;
quote! {
/// A serving function to use with [::tarpc::server::InFlightRequest::execute].
#[derive(Clone)]
#vis struct #server_ident<S> {
service: S,
}
}
}
fn impl_serve_for_server(&self) -> TokenStream2 {
let &Self {
request_ident,
server_ident,
service_ident,
response_ident,
camel_case_idents,
arg_pats,
method_idents,
request_names,
..
} = self;
quote! {
impl<S> ::tarpc::server::Serve for #server_ident<S>
where S: #service_ident
{
type Req = #request_ident;
type Resp = #response_ident;
fn method(&self, req: &#request_ident) -> ::core::option::Option<&'static str> {
::core::option::Option::Some(match req {
#(
#response_fut_ident::#camel_case_idents(resp) =>
std::pin::Pin::new_unchecked(resp)
.poll(cx)
.map(#response_ident::#camel_case_idents),
#request_ident::#camel_case_idents{..} => {
#request_names
}
)*
})
}
async fn serve(self, ctx: ::tarpc::context::Context, req: #request_ident)
-> ::core::result::Result<#response_ident, ::tarpc::ServerError> {
match req {
#(
#request_ident::#camel_case_idents{ #( #arg_pats ),* } => {
::core::result::Result::Ok(#response_ident::#camel_case_idents(
#service_ident::#method_idents(
self.service, ctx, #( #arg_pats ),*
).await
))
}
)*
}
}
}
}
}
}
fn struct_client(
&GenArgs {
vis,
client_ident,
request_ident,
response_ident,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
#[allow(unused)]
#[derive(Clone, Debug)]
/// The client stub that makes RPC calls to the server. Exposes a Future interface.
#vis struct #client_ident<C = tarpc::client::Channel<#request_ident, #response_ident>>(C);
}
}
fn enum_request(&self) -> TokenStream2 {
let &Self {
derive_serialize,
vis,
request_ident,
camel_case_idents,
args,
..
} = self;
fn impl_from_for_client(
&GenArgs {
client_ident,
request_ident,
response_ident,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
impl<C> From<C> for #client_ident<C>
where for <'a> C: tarpc::Client<'a, #request_ident, Response = #response_ident>
{
fn from(client: C) -> Self {
#client_ident(client)
quote! {
/// The request sent over the wire from the client to the server.
#[allow(missing_docs)]
#[derive(Debug)]
#derive_serialize
#vis enum #request_ident {
#( #camel_case_idents{ #( #args ),* } ),*
}
}
}
}
fn impl_client_new(
&GenArgs {
client_ident,
vis,
request_ident,
response_ident,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
impl #client_ident {
/// Returns a new client stub that sends requests over the given transport.
#vis fn new<T>(config: tarpc::client::Config, transport: T)
-> tarpc::client::NewClient<
Self,
tarpc::client::channel::RequestDispatch<#request_ident, #response_ident, T>>
where
T: tarpc::Transport<tarpc::ClientMessage<#request_ident>, tarpc::Response<#response_ident>>
{
let new_client = tarpc::client::new(config, transport);
tarpc::client::NewClient {
client: #client_ident(new_client.client),
dispatch: new_client.dispatch,
}
fn enum_response(&self) -> TokenStream2 {
let &Self {
derive_serialize,
vis,
response_ident,
camel_case_idents,
return_types,
..
} = self;
quote! {
/// The response sent over the wire from the server to the client.
#[allow(missing_docs)]
#[derive(Debug)]
#derive_serialize
#vis enum #response_ident {
#( #camel_case_idents(#return_types) ),*
}
}
}
}
fn impl_client_rpc_methods(
&GenArgs {
client_ident,
request_ident,
response_ident,
method_attrs,
vis,
method_names,
args,
return_types,
arg_vars,
camel_case_idents,
..
}: &GenArgs,
) -> proc_macro2::TokenStream {
quote! {
impl<C> #client_ident<C>
where for<'a> C: tarpc::Client<'a, #request_ident, Response = #response_ident>
{
#(
#[allow(unused)]
#( #method_attrs )*
#vis fn #method_names(&mut self, ctx: tarpc::context::Context, #( #args ),*)
-> impl std::future::Future<Output = std::io::Result<#return_types>> + '_ {
let request = #request_ident::#camel_case_idents { #( #arg_vars ),* };
let resp = tarpc::Client::call(&mut self.0, ctx, request);
async move {
match resp.await? {
#response_ident::#camel_case_idents(msg) => std::result::Result::Ok(msg),
_ => unreachable!(),
}
fn struct_client(&self) -> TokenStream2 {
let &Self {
vis,
client_ident,
request_ident,
response_ident,
..
} = self;
quote! {
#[allow(unused)]
#[derive(Clone, Debug)]
/// The client stub that makes RPC calls to the server. All request methods return
/// [Futures](::core::future::Future).
#vis struct #client_ident<
Stub = ::tarpc::client::Channel<#request_ident, #response_ident>
>(Stub);
}
}
fn impl_client_new(&self) -> TokenStream2 {
let &Self {
client_ident,
vis,
request_ident,
response_ident,
..
} = self;
quote! {
impl #client_ident {
/// Returns a new client stub that sends requests over the given transport.
#vis fn new<T>(config: ::tarpc::client::Config, transport: T)
-> ::tarpc::client::NewClient<
Self,
::tarpc::client::RequestDispatch<#request_ident, #response_ident, T>
>
where
T: ::tarpc::Transport<::tarpc::ClientMessage<#request_ident>, ::tarpc::Response<#response_ident>>
{
let new_client = ::tarpc::client::new(config, transport);
::tarpc::client::NewClient {
client: #client_ident(new_client.client),
dispatch: new_client.dispatch,
}
}
)*
}
impl<Stub> ::core::convert::From<Stub> for #client_ident<Stub>
where Stub: ::tarpc::client::stub::Stub<
Req = #request_ident,
Resp = #response_ident>
{
/// Returns a new client stub that sends requests over the given transport.
fn from(stub: Stub) -> Self {
#client_ident(stub)
}
}
}
}
fn impl_client_rpc_methods(&self) -> TokenStream2 {
let &Self {
client_ident,
request_ident,
response_ident,
method_attrs,
vis,
method_idents,
request_names,
args,
return_types,
arg_pats,
camel_case_idents,
..
} = self;
quote! {
impl<Stub> #client_ident<Stub>
where Stub: ::tarpc::client::stub::Stub<
Req = #request_ident,
Resp = #response_ident>
{
#(
#[allow(unused)]
#( #method_attrs )*
#vis fn #method_idents(&self, ctx: ::tarpc::context::Context, #( #args ),*)
-> impl ::core::future::Future<Output = ::core::result::Result<#return_types, ::tarpc::client::RpcError>> + '_ {
let request = #request_ident::#camel_case_idents { #( #arg_pats ),* };
let resp = self.0.call(ctx, #request_names, request);
async move {
match resp.await? {
#response_ident::#camel_case_idents(msg) => ::core::result::Result::Ok(msg),
_ => ::core::unreachable!(),
}
}
}
)*
}
}
}
}
impl<'a> ToTokens for ServiceGenerator<'a> {
fn to_tokens(&self, output: &mut TokenStream2) {
output.extend(vec![
self.trait_service(),
self.struct_server(),
self.impl_serve_for_server(),
self.enum_request(),
self.enum_response(),
self.struct_client(),
self.impl_client_new(),
self.impl_client_rpc_methods(),
])
}
}
fn snake_to_camel(ident_str: &str) -> String {
let chars = ident_str.chars();
let mut camel_ty = String::with_capacity(ident_str.len());
let mut last_char_was_underscore = true;
for c in chars {
for c in ident_str.chars() {
match c {
'_' => last_char_was_underscore = true,
c if last_char_was_underscore => {

44
plugins/tests/server.rs Normal file
View File

@@ -0,0 +1,44 @@
// these need to be out here rather than inside the function so that the
// assert_type_eq macro can pick them up.
#[tarpc::service]
trait Foo {
async fn two_part(s: String, i: i32) -> (String, i32);
async fn bar(s: String) -> String;
async fn baz();
}
#[allow(non_camel_case_types)]
#[test]
fn raw_idents_work() {
type r#yield = String;
#[tarpc::service]
trait r#trait {
async fn r#await(r#struct: r#yield, r#enum: i32) -> (r#yield, i32);
async fn r#fn(r#impl: r#yield) -> r#yield;
async fn r#async();
}
}
#[test]
fn syntax() {
#[tarpc::service]
trait Syntax {
#[deny(warnings)]
#[allow(non_snake_case)]
async fn TestCamelCaseDoesntConflict();
async fn hello() -> String;
#[doc = "attr"]
async fn attr(s: String) -> String;
async fn no_args_no_return();
async fn no_args() -> ();
async fn one_arg(one: String) -> i32;
async fn two_args_no_return(one: String, two: u64);
async fn two_args(one: String, two: u64) -> String;
async fn no_args_ret_error() -> i32;
async fn one_arg_ret_error(one: String) -> String;
async fn no_arg_implicit_return_error();
#[doc = "attr"]
async fn one_arg_implicit_return_error(one: String);
}
}

View File

@@ -2,8 +2,6 @@ use tarpc::context;
#[test]
fn att_service_trait() {
use futures::future::{ready, Ready};
#[tarpc::service]
trait Foo {
async fn two_part(s: String, i: i32) -> (String, i32);
@@ -12,19 +10,48 @@ fn att_service_trait() {
}
impl Foo for () {
type TwoPartFut = Ready<(String, i32)>;
fn two_part(self, _: context::Context, s: String, i: i32) -> Self::TwoPartFut {
ready((s, i))
async fn two_part(self, _: context::Context, s: String, i: i32) -> (String, i32) {
(s, i)
}
type BarFut = Ready<String>;
fn bar(self, _: context::Context, s: String) -> Self::BarFut {
ready(s)
async fn bar(self, _: context::Context, s: String) -> String {
s
}
type BazFut = Ready<()>;
fn baz(self, _: context::Context) -> Self::BazFut {
ready(())
async fn baz(self, _: context::Context) {
()
}
}
}
#[allow(non_camel_case_types)]
#[test]
fn raw_idents() {
type r#yield = String;
#[tarpc::service]
trait r#trait {
async fn r#await(r#struct: r#yield, r#enum: i32) -> (r#yield, i32);
async fn r#fn(r#impl: r#yield) -> r#yield;
async fn r#async();
}
impl r#trait for () {
async fn r#await(
self,
_: context::Context,
r#struct: r#yield,
r#enum: i32,
) -> (r#yield, i32) {
(r#struct, r#enum)
}
async fn r#fn(self, _: context::Context, r#impl: r#yield) -> r#yield {
r#impl
}
async fn r#async(self, _: context::Context) {
()
}
}
}

View File

@@ -1,57 +1,97 @@
[package]
name = "tarpc"
version = "0.20.0"
authors = ["Adam Wright <adam.austin.wright@gmail.com>", "Tim Kuehn <timothy.j.kuehn@gmail.com>"]
edition = "2018"
version = "0.34.0"
rust-version = "1.58.0"
authors = [
"Adam Wright <adam.austin.wright@gmail.com>",
"Tim Kuehn <timothy.j.kuehn@gmail.com>",
]
edition = "2021"
license = "MIT"
documentation = "https://docs.rs/tarpc"
homepage = "https://github.com/google/tarpc"
repository = "https://github.com/google/tarpc"
keywords = ["rpc", "network", "server", "api", "microservices"]
categories = ["asynchronous", "network-programming"]
readme = "../README.md"
readme = "README.md"
description = "An RPC framework for Rust with a focus on ease of use."
[features]
default = []
serde1 = ["tarpc-plugins/serde1", "serde", "serde/derive"]
tokio1 = []
serde-transport = ["tokio-serde", "tokio-util/codec"]
tcp = ["tokio/net", "tokio/stream"]
serde1 = ["tarpc-plugins/serde1", "serde", "serde/derive", "serde/rc"]
tokio1 = ["tokio/rt"]
serde-transport = ["serde1", "tokio1", "tokio-serde", "tokio-util/codec"]
serde-transport-json = ["tokio-serde/json"]
serde-transport-bincode = ["tokio-serde/bincode"]
tcp = ["tokio/net"]
unix = ["tokio/net"]
full = ["serde1", "tokio1", "serde-transport", "tcp"]
full = [
"serde1",
"tokio1",
"serde-transport",
"serde-transport-json",
"serde-transport-bincode",
"tcp",
"unix",
]
[badges]
travis-ci = { repository = "google/tarpc" }
[dependencies]
anyhow = "1.0"
fnv = "1.0"
futures = "0.3"
humantime = "1.0"
log = "0.4"
pin-project = "0.4"
raii-counter = "0.2"
rand = "0.7"
tokio = { version = "0.2", features = ["time"] }
humantime = "2.0"
pin-project = "1.0"
rand = "0.8"
serde = { optional = true, version = "1.0", features = ["derive"] }
tokio-util = { optional = true, version = "0.2" }
tarpc-plugins = { path = "../plugins", version = "0.7" }
tokio-serde = { optional = true, version = "0.6" }
static_assertions = "1.1.0"
tarpc-plugins = { path = "../plugins", version = "0.13" }
thiserror = "1.0"
tokio = { version = "1", features = ["time"] }
tokio-util = { version = "0.7.3", features = ["time"] }
tokio-serde = { optional = true, version = "0.8" }
tracing = { version = "0.1", default-features = false, features = [
"attributes",
"log",
] }
tracing-opentelemetry = { version = "0.18.0", default-features = false }
opentelemetry = { version = "0.18.0", default-features = false }
[dev-dependencies]
assert_matches = "1.0"
bytes = { version = "0.5", features = ["serde"] }
env_logger = "0.6"
futures = "0.3"
humantime = "1.0"
log = "0.4"
assert_matches = "1.4"
bincode = "1.3"
bytes = { version = "1", features = ["serde"] }
flate2 = "1.0"
futures-test = "0.3"
opentelemetry = { version = "0.18.0", default-features = false, features = [
"rt-tokio",
] }
opentelemetry-jaeger = { version = "0.17.0", features = ["rt-tokio"] }
pin-utils = "0.1.0-alpha"
tokio = { version = "0.2", features = ["full"] }
tokio-serde = { version = "0.6", features = ["json"] }
serde_bytes = "0.11"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
tokio = { version = "1", features = ["full", "test-util", "tracing"] }
console-subscriber = "0.1"
tokio-serde = { version = "0.8", features = ["json", "bincode"] }
trybuild = "1.0"
tokio-rustls = "0.23"
rustls-pemfile = "1.0"
[package.metadata.docs.rs]
all-features = true
rustdoc-args = ["--cfg", "docsrs"]
[[example]]
name = "server_calling_server"
name = "compression"
required-features = ["serde-transport", "tcp"]
[[example]]
name = "tracing"
required-features = ["full"]
[[example]]
@@ -62,3 +102,18 @@ required-features = ["full"]
name = "pubsub"
required-features = ["full"]
[[example]]
name = "custom_transport"
required-features = ["serde1", "tokio1", "serde-transport"]
[[example]]
name = "tls_over_tcp"
required-features = ["full"]
[[test]]
name = "service_functional"
required-features = ["serde-transport"]
[[test]]
name = "dataservice"
required-features = ["serde-transport", "tcp"]

9
tarpc/LICENSE Normal file
View File

@@ -0,0 +1,9 @@
The MIT License (MIT)
Copyright 2016 Google Inc. All Rights Reserved.
Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

View File

@@ -0,0 +1,11 @@
-----BEGIN CERTIFICATE-----
MIIBlDCCAUagAwIBAgICAxUwBQYDK2VwMC4xLDAqBgNVBAMMI3Bvbnl0b3duIEVk
RFNBIGxldmVsIDIgaW50ZXJtZWRpYXRlMB4XDTIzMDMxNzEwMTEwNFoXDTI4MDkw
NjEwMTEwNFowGjEYMBYGA1UEAwwPcG9ueXRvd24gY2xpZW50MCowBQYDK2VwAyEA
NTKuLume19IhJfEFd/5OZUuYDKZH6xvy4AGver17OoejgZswgZgwDAYDVR0TAQH/
BAIwADALBgNVHQ8EBAMCBsAwFgYDVR0lAQH/BAwwCgYIKwYBBQUHAwIwHQYDVR0O
BBYEFDjdrlMu4tyw5MHtbg7WnzSGRBpFMEQGA1UdIwQ9MDuAFHIl7fHKWP6/l8FE
fI2YEIM3oHxKoSCkHjAcMRowGAYDVQQDDBFwb255dG93biBFZERTQSBDQYIBezAF
BgMrZXADQQCaahfj/QLxoCOpvl6y0ZQ9CpojPqBnxV3460j5nUOp040Va2MpF137
izCBY7LwgUE/YG6E+kH30G4jMEnqVEYK
-----END CERTIFICATE-----

View File

@@ -0,0 +1,19 @@
-----BEGIN CERTIFICATE-----
MIIBeDCCASqgAwIBAgIBezAFBgMrZXAwHDEaMBgGA1UEAwwRcG9ueXRvd24gRWRE
U0EgQ0EwHhcNMjMwMzE3MTAxMTA0WhcNMzMwMzE0MTAxMTA0WjAuMSwwKgYDVQQD
DCNwb255dG93biBFZERTQSBsZXZlbCAyIGludGVybWVkaWF0ZTAqMAUGAytlcAMh
AEFsAexz4x2R4k4+PnTbvRVn0r3F/qw/zVnNBxfGcoEpo38wfTAdBgNVHQ4EFgQU
ciXt8cpY/r+XwUR8jZgQgzegfEowIAYDVR0lAQH/BBYwFAYIKwYBBQUHAwEGCCsG
AQUFBwMCMAwGA1UdEwQFMAMBAf8wCwYDVR0PBAQDAgH+MB8GA1UdIwQYMBaAFKYU
oLdKeY7mp7QgMZKrkVtSWYBKMAUGAytlcANBAHVpNpCV8nu4fkH3Smikx5A9qtHc
zgLIyp+wrF1a4YSa6sfTvuQmJd5aF23OXgq5grCOPXtdpHO50Mx5Qy74zQg=
-----END CERTIFICATE-----
-----BEGIN CERTIFICATE-----
MIIBTDCB/6ADAgECAhRZLuF0TWjDs/31OO8VeKHkNIJQaDAFBgMrZXAwHDEaMBgG
A1UEAwwRcG9ueXRvd24gRWREU0EgQ0EwHhcNMjMwMzE3MTAxMTA0WhcNMzMwMzE0
MTAxMTA0WjAcMRowGAYDVQQDDBFwb255dG93biBFZERTQSBDQTAqMAUGAytlcAMh
ABRPZ4TiuBE8CqAFByZvqpMo/unjnnryfG2AkkWGXpa3o1MwUTAdBgNVHQ4EFgQU
phSgt0p5juantCAxkquRW1JZgEowHwYDVR0jBBgwFoAUphSgt0p5juantCAxkquR
W1JZgEowDwYDVR0TAQH/BAUwAwEB/zAFBgMrZXADQQB29o8erJA0/a8/xOHilOCC
t/s5wPHHnS5NSKx/m2N2nRn3zPxEnETlrAmGulJoeKOx8OblwmPi9rBT2K+QY2UB
-----END CERTIFICATE-----

View File

@@ -0,0 +1,3 @@
-----BEGIN PRIVATE KEY-----
MC4CAQAwBQYDK2VwBCIEIIJX9ThTHpVS1SNZb6HP4myg4fRInIVGunTRdgnc+weH
-----END PRIVATE KEY-----

View File

@@ -0,0 +1,12 @@
-----BEGIN CERTIFICATE-----
MIIBuDCCAWqgAwIBAgICAcgwBQYDK2VwMC4xLDAqBgNVBAMMI3Bvbnl0b3duIEVk
RFNBIGxldmVsIDIgaW50ZXJtZWRpYXRlMB4XDTIzMDMxNzEwMTEwNFoXDTI4MDkw
NjEwMTEwNFowGTEXMBUGA1UEAwwOdGVzdHNlcnZlci5jb20wKjAFBgMrZXADIQDc
RLl3/N2tPoWnzBV3noVn/oheEl8IUtiY11Vg/QXTUKOBwDCBvTAMBgNVHRMBAf8E
AjAAMAsGA1UdDwQEAwIGwDAdBgNVHQ4EFgQUk7U2mnxedNWBAH84BsNy5si3ZQow
RAYDVR0jBD0wO4AUciXt8cpY/r+XwUR8jZgQgzegfEqhIKQeMBwxGjAYBgNVBAMM
EXBvbnl0b3duIEVkRFNBIENBggF7MDsGA1UdEQQ0MDKCDnRlc3RzZXJ2ZXIuY29t
ghVzZWNvbmQudGVzdHNlcnZlci5jb22CCWxvY2FsaG9zdDAFBgMrZXADQQCFWIcF
9FiztCuUNzgXDNu5kshuflt0RjkjWpGlWzQjGoYM2IvYhNVPeqnCiY92gqwDSBtq
amD2TBup4eNUCsQB
-----END CERTIFICATE-----

View File

@@ -0,0 +1,19 @@
-----BEGIN CERTIFICATE-----
MIIBeDCCASqgAwIBAgIBezAFBgMrZXAwHDEaMBgGA1UEAwwRcG9ueXRvd24gRWRE
U0EgQ0EwHhcNMjMwMzE3MTAxMTA0WhcNMzMwMzE0MTAxMTA0WjAuMSwwKgYDVQQD
DCNwb255dG93biBFZERTQSBsZXZlbCAyIGludGVybWVkaWF0ZTAqMAUGAytlcAMh
AEFsAexz4x2R4k4+PnTbvRVn0r3F/qw/zVnNBxfGcoEpo38wfTAdBgNVHQ4EFgQU
ciXt8cpY/r+XwUR8jZgQgzegfEowIAYDVR0lAQH/BBYwFAYIKwYBBQUHAwEGCCsG
AQUFBwMCMAwGA1UdEwQFMAMBAf8wCwYDVR0PBAQDAgH+MB8GA1UdIwQYMBaAFKYU
oLdKeY7mp7QgMZKrkVtSWYBKMAUGAytlcANBAHVpNpCV8nu4fkH3Smikx5A9qtHc
zgLIyp+wrF1a4YSa6sfTvuQmJd5aF23OXgq5grCOPXtdpHO50Mx5Qy74zQg=
-----END CERTIFICATE-----
-----BEGIN CERTIFICATE-----
MIIBTDCB/6ADAgECAhRZLuF0TWjDs/31OO8VeKHkNIJQaDAFBgMrZXAwHDEaMBgG
A1UEAwwRcG9ueXRvd24gRWREU0EgQ0EwHhcNMjMwMzE3MTAxMTA0WhcNMzMwMzE0
MTAxMTA0WjAcMRowGAYDVQQDDBFwb255dG93biBFZERTQSBDQTAqMAUGAytlcAMh
ABRPZ4TiuBE8CqAFByZvqpMo/unjnnryfG2AkkWGXpa3o1MwUTAdBgNVHQ4EFgQU
phSgt0p5juantCAxkquRW1JZgEowHwYDVR0jBBgwFoAUphSgt0p5juantCAxkquR
W1JZgEowDwYDVR0TAQH/BAUwAwEB/zAFBgMrZXADQQB29o8erJA0/a8/xOHilOCC
t/s5wPHHnS5NSKx/m2N2nRn3zPxEnETlrAmGulJoeKOx8OblwmPi9rBT2K+QY2UB
-----END CERTIFICATE-----

View File

@@ -0,0 +1,3 @@
-----BEGIN PRIVATE KEY-----
MC4CAQAwBQYDK2VwBCIEIMU6xGVe8JTpZ3bN/wajHfw6pEHt0Rd7wPBxds9eEFy2
-----END PRIVATE KEY-----

View File

@@ -0,0 +1,138 @@
// Copyright 2022 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use flate2::{read::DeflateDecoder, write::DeflateEncoder, Compression};
use futures::{prelude::*, Sink, SinkExt, Stream, StreamExt, TryStreamExt};
use serde::{Deserialize, Serialize};
use serde_bytes::ByteBuf;
use std::{io, io::Read, io::Write};
use tarpc::{
client, context,
serde_transport::tcp,
server::{BaseChannel, Channel},
tokio_serde::formats::Bincode,
};
/// Type of compression that should be enabled on the request. The transport is free to ignore this.
#[derive(Debug, PartialEq, Eq, Clone, Copy, Deserialize, Serialize)]
pub enum CompressionAlgorithm {
Deflate,
}
#[derive(Debug, Deserialize, Serialize)]
pub enum CompressedMessage<T> {
Uncompressed(T),
Compressed {
algorithm: CompressionAlgorithm,
payload: ByteBuf,
},
}
#[derive(Deserialize, Serialize)]
enum CompressionType {
Uncompressed,
Compressed,
}
async fn compress<T>(message: T) -> io::Result<CompressedMessage<T>>
where
T: Serialize,
{
let message = serialize(message)?;
let mut encoder = DeflateEncoder::new(Vec::new(), Compression::default());
encoder.write_all(&message).unwrap();
let compressed = encoder.finish()?;
Ok(CompressedMessage::Compressed {
algorithm: CompressionAlgorithm::Deflate,
payload: ByteBuf::from(compressed),
})
}
async fn decompress<T>(message: CompressedMessage<T>) -> io::Result<T>
where
for<'a> T: Deserialize<'a>,
{
match message {
CompressedMessage::Compressed { algorithm, payload } => {
if algorithm != CompressionAlgorithm::Deflate {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!("Compression algorithm {algorithm:?} not supported"),
));
}
let mut deflater = DeflateDecoder::new(payload.as_slice());
let mut payload = ByteBuf::new();
deflater.read_to_end(&mut payload)?;
let message = deserialize(payload)?;
Ok(message)
}
CompressedMessage::Uncompressed(message) => Ok(message),
}
}
fn serialize<T: Serialize>(t: T) -> io::Result<ByteBuf> {
bincode::serialize(&t)
.map(ByteBuf::from)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
fn deserialize<D>(message: ByteBuf) -> io::Result<D>
where
for<'a> D: Deserialize<'a>,
{
bincode::deserialize(message.as_ref()).map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
fn add_compression<In, Out>(
transport: impl Stream<Item = io::Result<CompressedMessage<In>>>
+ Sink<CompressedMessage<Out>, Error = io::Error>,
) -> impl Stream<Item = io::Result<In>> + Sink<Out, Error = io::Error>
where
Out: Serialize,
for<'a> In: Deserialize<'a>,
{
transport.with(compress).and_then(decompress)
}
#[tarpc::service]
pub trait World {
async fn hello(name: String) -> String;
}
#[derive(Clone, Debug)]
struct HelloServer;
impl World for HelloServer {
async fn hello(self, _: context::Context, name: String) -> String {
format!("Hey, {name}!")
}
}
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let mut incoming = tcp::listen("localhost:0", Bincode::default).await?;
let addr = incoming.local_addr();
tokio::spawn(async move {
let transport = incoming.next().await.unwrap().unwrap();
BaseChannel::with_defaults(add_compression(transport))
.execute(HelloServer.serve())
.for_each(spawn)
.await;
});
let transport = tcp::connect(addr, Bincode::default).await?;
let client = WorldClient::new(client::Config::default(), add_compression(transport)).spawn();
println!(
"{}",
client.hello(context::current(), "friend".into()).await?
);
Ok(())
}

View File

@@ -0,0 +1,59 @@
// Copyright 2022 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use futures::prelude::*;
use tarpc::context::Context;
use tarpc::serde_transport as transport;
use tarpc::server::{BaseChannel, Channel};
use tarpc::tokio_serde::formats::Bincode;
use tarpc::tokio_util::codec::length_delimited::LengthDelimitedCodec;
use tokio::net::{UnixListener, UnixStream};
#[tarpc::service]
pub trait PingService {
async fn ping();
}
#[derive(Clone)]
struct Service;
impl PingService for Service {
async fn ping(self, _: Context) {}
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let bind_addr = "/tmp/tarpc_on_unix_example.sock";
let _ = std::fs::remove_file(bind_addr);
let listener = UnixListener::bind(bind_addr).unwrap();
let codec_builder = LengthDelimitedCodec::builder();
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
tokio::spawn(async move {
loop {
let (conn, _addr) = listener.accept().await.unwrap();
let framed = codec_builder.new_framed(conn);
let transport = transport::new(framed, Bincode::default());
let fut = BaseChannel::with_defaults(transport)
.execute(Service.serve())
.for_each(spawn);
tokio::spawn(fut);
}
});
let conn = UnixStream::connect(bind_addr).await?;
let transport = transport::new(codec_builder.new_framed(conn), Bincode::default());
PingServiceClient::new(Default::default(), transport)
.spawn()
.ping(tarpc::context::current())
.await?;
Ok(())
}

View File

@@ -4,192 +4,359 @@
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
/// - The PubSub server sets up TCP listeners on 2 ports, the "subscriber" port and the "publisher"
/// port. Because both publishers and subscribers initiate their connections to the PubSub
/// server, the server requires no prior knowledge of either publishers or subscribers.
///
/// - Subscribers connect to the server on the server's "subscriber" port. Once a connection is
/// established, the server acts as the client of the Subscriber service, initially requesting
/// the topics the subscriber is interested in, and subsequently sending topical messages to the
/// subscriber.
///
/// - Publishers connect to the server on the "publisher" port and, once connected, they send
/// topical messages via Publisher service to the server. The server then broadcasts each
/// messages to all clients subscribed to the topic of that message.
///
/// Subscriber Publisher PubSub Server
/// T1 | | |
/// T2 |-----Connect------------------------------------------------------>|
/// T3 | | |
/// T2 |<-------------------------------------------------------Topics-----|
/// T2 |-----(OK) Topics-------------------------------------------------->|
/// T3 | | |
/// T4 | |-----Connect-------------------->|
/// T5 | | |
/// T6 | |-----Publish-------------------->|
/// T7 | | |
/// T8 |<------------------------------------------------------Receive-----|
/// T9 |-----(OK) Receive------------------------------------------------->|
/// T10 | | |
/// T11 | |<--------------(OK) Publish------|
use anyhow::anyhow;
use futures::{
future::{self, Ready},
channel::oneshot,
future::{self, AbortHandle},
prelude::*,
Future,
};
use publisher::Publisher as _;
use std::{
collections::HashMap,
env,
error::Error,
io,
net::SocketAddr,
pin::Pin,
sync::{Arc, Mutex},
time::Duration,
sync::{Arc, Mutex, RwLock},
};
use subscriber::Subscriber as _;
use tarpc::{
client, context,
server::{self, Handler},
serde_transport::tcp,
server::{self, Channel},
tokio_serde::formats::Json,
};
use tokio_serde::formats::Json;
use tokio::net::ToSocketAddrs;
use tracing::info;
use tracing_subscriber::prelude::*;
pub mod subscriber {
#[tarpc::service]
pub trait Subscriber {
async fn receive(message: String);
async fn topics() -> Vec<String>;
async fn receive(topic: String, message: String);
}
}
pub mod publisher {
use std::net::SocketAddr;
#[tarpc::service]
pub trait Publisher {
async fn broadcast(message: String);
async fn subscribe(id: u32, address: SocketAddr) -> Result<(), String>;
async fn unsubscribe(id: u32);
async fn publish(topic: String, message: String);
}
}
#[derive(Clone, Debug)]
struct Subscriber {
id: u32,
local_addr: SocketAddr,
topics: Vec<String>,
}
impl subscriber::Subscriber for Subscriber {
type ReceiveFut = Ready<()>;
async fn topics(self, _: context::Context) -> Vec<String> {
self.topics.clone()
}
fn receive(self, _: context::Context, message: String) -> Self::ReceiveFut {
eprintln!("{} received message: {}", self.id, message);
future::ready(())
async fn receive(self, _: context::Context, topic: String, message: String) {
info!(local_addr = %self.local_addr, %topic, %message, "ReceivedMessage")
}
}
struct SubscriberHandle(AbortHandle);
impl Drop for SubscriberHandle {
fn drop(&mut self) {
self.0.abort();
}
}
impl Subscriber {
async fn listen(id: u32, config: server::Config) -> io::Result<SocketAddr> {
let incoming = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()));
let addr = incoming.get_ref().local_addr();
tokio::spawn(
server::new(config)
.incoming(incoming)
.take(1)
.respond_with(Subscriber { id }.serve()),
);
Ok(addr)
async fn connect(
publisher_addr: impl ToSocketAddrs,
topics: Vec<String>,
) -> anyhow::Result<SubscriberHandle> {
let publisher = tcp::connect(publisher_addr, Json::default).await?;
let local_addr = publisher.local_addr()?;
let mut handler = server::BaseChannel::with_defaults(publisher).requests();
let subscriber = Subscriber { local_addr, topics };
// The first request is for the topics being subscribed to.
match handler.next().await {
Some(init_topics) => init_topics?.execute(subscriber.clone().serve()).await,
None => {
return Err(anyhow!(
"[{}] Server never initialized the subscriber.",
local_addr
))
}
};
let (handler, abort_handle) =
future::abortable(handler.execute(subscriber.serve()).for_each(spawn));
tokio::spawn(async move {
match handler.await {
Ok(()) | Err(future::Aborted) => info!(?local_addr, "subscriber shutdown."),
}
});
Ok(SubscriberHandle(abort_handle))
}
}
#[derive(Debug)]
struct Subscription {
topics: Vec<String>,
}
#[derive(Clone, Debug)]
struct Publisher {
clients: Arc<Mutex<HashMap<u32, subscriber::SubscriberClient>>>,
clients: Arc<Mutex<HashMap<SocketAddr, Subscription>>>,
subscriptions: Arc<RwLock<HashMap<String, HashMap<SocketAddr, subscriber::SubscriberClient>>>>,
}
struct PublisherAddrs {
publisher: SocketAddr,
subscriptions: SocketAddr,
}
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
impl Publisher {
fn new() -> Publisher {
Publisher {
clients: Arc::new(Mutex::new(HashMap::new())),
async fn start(self) -> io::Result<PublisherAddrs> {
let mut connecting_publishers = tcp::listen("localhost:0", Json::default).await?;
let publisher_addrs = PublisherAddrs {
publisher: connecting_publishers.local_addr(),
subscriptions: self.clone().start_subscription_manager().await?,
};
info!(publisher_addr = %publisher_addrs.publisher, "listening for publishers.",);
tokio::spawn(async move {
// Because this is just an example, we know there will only be one publisher. In more
// realistic code, this would be a loop to continually accept new publisher
// connections.
let publisher = connecting_publishers.next().await.unwrap().unwrap();
info!(publisher.peer_addr = ?publisher.peer_addr(), "publisher connected.");
server::BaseChannel::with_defaults(publisher)
.execute(self.serve())
.for_each(spawn)
.await
});
Ok(publisher_addrs)
}
async fn start_subscription_manager(mut self) -> io::Result<SocketAddr> {
let mut connecting_subscribers = tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()));
let new_subscriber_addr = connecting_subscribers.get_ref().local_addr();
info!(?new_subscriber_addr, "listening for subscribers.");
tokio::spawn(async move {
while let Some(conn) = connecting_subscribers.next().await {
let subscriber_addr = conn.peer_addr().unwrap();
let tarpc::client::NewClient {
client: subscriber,
dispatch,
} = subscriber::SubscriberClient::new(client::Config::default(), conn);
let (ready_tx, ready) = oneshot::channel();
self.clone()
.start_subscriber_gc(subscriber_addr, dispatch, ready);
// Populate the topics
self.initialize_subscription(subscriber_addr, subscriber)
.await;
// Signal that initialization is done.
ready_tx.send(()).unwrap();
}
});
Ok(new_subscriber_addr)
}
async fn initialize_subscription(
&mut self,
subscriber_addr: SocketAddr,
subscriber: subscriber::SubscriberClient,
) {
// Populate the topics
if let Ok(topics) = subscriber.topics(context::current()).await {
self.clients.lock().unwrap().insert(
subscriber_addr,
Subscription {
topics: topics.clone(),
},
);
info!(%subscriber_addr, ?topics, "subscribed to new topics");
let mut subscriptions = self.subscriptions.write().unwrap();
for topic in topics {
subscriptions
.entry(topic)
.or_insert_with(HashMap::new)
.insert(subscriber_addr, subscriber.clone());
}
}
}
fn start_subscriber_gc<E: Error>(
self,
subscriber_addr: SocketAddr,
client_dispatch: impl Future<Output = Result<(), E>> + Send + 'static,
subscriber_ready: oneshot::Receiver<()>,
) {
tokio::spawn(async move {
if let Err(e) = client_dispatch.await {
info!(
%subscriber_addr,
error = %e,
"subscriber connection broken");
}
// Don't clean up the subscriber until initialization is done.
let _ = subscriber_ready.await;
if let Some(subscription) = self.clients.lock().unwrap().remove(&subscriber_addr) {
info!(
"[{} unsubscribing from topics: {:?}",
subscriber_addr, subscription.topics
);
let mut subscriptions = self.subscriptions.write().unwrap();
for topic in subscription.topics {
let subscribers = subscriptions.get_mut(&topic).unwrap();
subscribers.remove(&subscriber_addr);
if subscribers.is_empty() {
subscriptions.remove(&topic);
}
}
}
});
}
}
impl publisher::Publisher for Publisher {
type BroadcastFut = Pin<Box<dyn Future<Output = ()> + Send>>;
fn broadcast(self, _: context::Context, message: String) -> Self::BroadcastFut {
async fn broadcast(
clients: Arc<Mutex<HashMap<u32, subscriber::SubscriberClient>>>,
message: String,
) {
let mut clients = clients.lock().unwrap().clone();
for client in clients.values_mut() {
// Ignore failing subscribers. In a real pubsub,
// you'd want to continually retry until subscribers
// ack.
let _ = client.receive(context::current(), message.clone()).await;
async fn publish(self, _: context::Context, topic: String, message: String) {
info!("received message to publish.");
let mut subscribers = match self.subscriptions.read().unwrap().get(&topic) {
None => return,
Some(subscriptions) => subscriptions.clone(),
};
let mut publications = Vec::new();
for client in subscribers.values_mut() {
publications.push(client.receive(context::current(), topic.clone(), message.clone()));
}
// Ignore failing subscribers. In a real pubsub, you'd want to continually retry until
// subscribers ack. Of course, a lot would be different in a real pubsub :)
for response in future::join_all(publications).await {
if let Err(e) = response {
info!("failed to broadcast to subscriber: {}", e);
}
}
broadcast(self.clients.clone(), message).boxed()
}
type SubscribeFut = Pin<Box<dyn Future<Output = Result<(), String>> + Send>>;
fn subscribe(self, _: context::Context, id: u32, addr: SocketAddr) -> Self::SubscribeFut {
async fn subscribe(
clients: Arc<Mutex<HashMap<u32, subscriber::SubscriberClient>>>,
id: u32,
addr: SocketAddr,
) -> io::Result<()> {
let conn = tarpc::serde_transport::tcp::connect(addr, Json::default()).await?;
let subscriber =
subscriber::SubscriberClient::new(client::Config::default(), conn).spawn()?;
eprintln!("Subscribing {}.", id);
clients.lock().unwrap().insert(id, subscriber);
Ok(())
}
subscribe(Arc::clone(&self.clients), id, addr)
.map_err(|e| e.to_string())
.boxed()
}
type UnsubscribeFut = Pin<Box<dyn Future<Output = ()> + Send>>;
fn unsubscribe(self, _: context::Context, id: u32) -> Self::UnsubscribeFut {
eprintln!("Unsubscribing {}", id);
let mut clients = self.clients.lock().unwrap();
if clients.remove(&id).is_none() {
eprintln!(
"Client {} not found. Existings clients: {:?}",
id, &*clients
);
}
future::ready(()).boxed()
}
}
#[tokio::main]
async fn main() -> io::Result<()> {
env_logger::init();
/// Initializes an OpenTelemetry tracing subscriber with a Jaeger backend.
fn init_tracing(service_name: &str) -> anyhow::Result<()> {
env::set_var("OTEL_BSP_MAX_EXPORT_BATCH_SIZE", "12");
let tracer = opentelemetry_jaeger::new_agent_pipeline()
.with_service_name(service_name)
.with_max_packet_size(2usize.pow(13))
.install_batch(opentelemetry::runtime::Tokio)?;
let transport = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()));
let publisher_addr = transport.get_ref().local_addr();
tokio::spawn(
transport
.take(1)
.map(server::BaseChannel::with_defaults)
.respond_with(Publisher::new().serve()),
);
let subscriber1 = Subscriber::listen(0, server::Config::default()).await?;
let subscriber2 = Subscriber::listen(1, server::Config::default()).await?;
let publisher_conn = tarpc::serde_transport::tcp::connect(publisher_addr, Json::default());
let publisher_conn = publisher_conn.await?;
let mut publisher =
publisher::PublisherClient::new(client::Config::default(), publisher_conn).spawn()?;
if let Err(e) = publisher
.subscribe(context::current(), 0, subscriber1)
.await?
{
eprintln!("Couldn't subscribe subscriber 0: {}", e);
}
if let Err(e) = publisher
.subscribe(context::current(), 1, subscriber2)
.await?
{
eprintln!("Couldn't subscribe subscriber 1: {}", e);
}
println!("Broadcasting...");
publisher
.broadcast(context::current(), "hello to all".to_string())
.await?;
publisher.unsubscribe(context::current(), 1).await?;
publisher
.broadcast(context::current(), "hi again".to_string())
.await?;
drop(publisher);
tokio::time::delay_for(Duration::from_millis(100)).await;
println!("Done.");
tracing_subscriber::registry()
.with(tracing_subscriber::filter::EnvFilter::from_default_env())
.with(tracing_subscriber::fmt::layer())
.with(tracing_opentelemetry::layer().with_tracer(tracer))
.try_init()?;
Ok(())
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
init_tracing("Pub/Sub")?;
let addrs = Publisher {
clients: Arc::new(Mutex::new(HashMap::new())),
subscriptions: Arc::new(RwLock::new(HashMap::new())),
}
.start()
.await?;
let _subscriber0 = Subscriber::connect(
addrs.subscriptions,
vec!["calculus".into(), "cool shorts".into()],
)
.await?;
let _subscriber1 = Subscriber::connect(
addrs.subscriptions,
vec!["cool shorts".into(), "history".into()],
)
.await?;
let publisher = publisher::PublisherClient::new(
client::Config::default(),
tcp::connect(addrs.publisher, Json::default).await?,
)
.spawn();
publisher
.publish(context::current(), "calculus".into(), "sqrt(2)".into())
.await?;
publisher
.publish(
context::current(),
"cool shorts".into(),
"hello to all".into(),
)
.await?;
publisher
.publish(context::current(), "history".into(), "napoleon".to_string())
.await?;
drop(_subscriber0);
publisher
.publish(
context::current(),
"cool shorts".into(),
"hello to who?".into(),
)
.await?;
opentelemetry::global::shutdown_tracer_provider();
info!("done.");
Ok(())
}

View File

@@ -4,16 +4,11 @@
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use futures::{
future::{self, Ready},
prelude::*,
};
use std::io;
use futures::prelude::*;
use tarpc::{
client, context,
server::{BaseChannel, Channel},
server::{self, Channel},
};
use tokio_serde::formats::Json;
/// This is the service definition. It looks a lot like a trait definition.
/// It defines one RPC, hello, which takes one arg, name, and returns a String.
@@ -28,52 +23,32 @@ pub trait World {
struct HelloServer;
impl World for HelloServer {
// Each defined rpc generates two items in the trait, a fn that serves the RPC, and
// an associated type representing the future output by the fn.
type HelloFut = Ready<String>;
fn hello(self, _: context::Context, name: String) -> Self::HelloFut {
future::ready(format!("Hello, {}!", name))
async fn hello(self, _: context::Context, name: String) -> String {
format!("Hello, {name}!")
}
}
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
#[tokio::main]
async fn main() -> io::Result<()> {
// tarpc_json_transport is provided by the associated crate json_transport. It makes it
// easy to start up a serde-powered JSON serialization strategy over TCP.
let mut transport = tarpc::serde_transport::tcp::listen("localhost:0", Json::default).await?;
let addr = transport.local_addr();
async fn main() -> anyhow::Result<()> {
let (client_transport, server_transport) = tarpc::transport::channel::unbounded();
let server = async move {
// For this example, we're just going to wait for one connection.
let client = transport.next().await.unwrap().unwrap();
let server = server::BaseChannel::with_defaults(server_transport);
tokio::spawn(server.execute(HelloServer.serve()).for_each(spawn));
// `Channel` is a trait representing a server-side connection. It is a trait to allow
// for some channels to be instrumented: for example, to track the number of open connections.
// BaseChannel is the most basic channel, simply wrapping a transport with no added
// functionality.
BaseChannel::with_defaults(client)
// serve_world is generated by the tarpc::service attribute. It takes as input any type
// implementing the generated World trait.
.respond_with(HelloServer.serve())
.execute()
.await;
};
tokio::spawn(server);
let transport = tarpc::serde_transport::tcp::connect(addr, Json::default()).await?;
// WorldClient is generated by the tarpc::service attribute. It has a constructor `new` that
// takes a config and any Transport as input.
let mut client = WorldClient::new(client::Config::default(), transport).spawn()?;
// WorldClient is generated by the #[tarpc::service] attribute. It has a constructor `new`
// that takes a config and any Transport as input.
let client = WorldClient::new(client::Config::default(), client_transport).spawn();
// The client has an RPC method for each RPC defined in the annotated trait. It takes the same
// args as defined, with the addition of a Context, which is always the first arg. The Context
// specifies a deadline and trace information which can be helpful in debugging requests.
let hello = client.hello(context::current(), "Stim".to_string()).await?;
eprintln!("{}", hello);
println!("{hello}");
Ok(())
}

View File

@@ -1,101 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use crate::{add::Add as AddService, double::Double as DoubleService};
use futures::{
future::{self, Ready},
prelude::*,
};
use std::{io, pin::Pin};
use tarpc::{
client, context,
server::{Handler, Server},
};
use tokio_serde::formats::Json;
pub mod add {
#[tarpc::service]
pub trait Add {
/// Add two ints together.
async fn add(x: i32, y: i32) -> i32;
}
}
pub mod double {
#[tarpc::service]
pub trait Double {
/// 2 * x
async fn double(x: i32) -> Result<i32, String>;
}
}
#[derive(Clone)]
struct AddServer;
impl AddService for AddServer {
type AddFut = Ready<i32>;
fn add(self, _: context::Context, x: i32, y: i32) -> Self::AddFut {
future::ready(x + y)
}
}
#[derive(Clone)]
struct DoubleServer {
add_client: add::AddClient,
}
impl DoubleService for DoubleServer {
type DoubleFut = Pin<Box<dyn Future<Output = Result<i32, String>> + Send>>;
fn double(self, _: context::Context, x: i32) -> Self::DoubleFut {
async fn double(mut client: add::AddClient, x: i32) -> Result<i32, String> {
client
.add(context::current(), x, x)
.await
.map_err(|e| e.to_string())
}
double(self.add_client.clone(), x).boxed()
}
}
#[tokio::main]
async fn main() -> io::Result<()> {
env_logger::init();
let add_listener = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()));
let addr = add_listener.get_ref().local_addr();
let add_server = Server::default()
.incoming(add_listener)
.take(1)
.respond_with(AddServer.serve());
tokio::spawn(add_server);
let to_add_server = tarpc::serde_transport::tcp::connect(addr, Json::default()).await?;
let add_client = add::AddClient::new(client::Config::default(), to_add_server).spawn()?;
let double_listener = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()));
let addr = double_listener.get_ref().local_addr();
let double_server = tarpc::Server::default()
.incoming(double_listener)
.take(1)
.respond_with(DoubleServer { add_client }.serve());
tokio::spawn(double_server);
let to_double_server = tarpc::serde_transport::tcp::connect(addr, Json::default()).await?;
let mut double_client =
double::DoubleClient::new(client::Config::default(), to_double_server).spawn()?;
for i in 1..=5 {
eprintln!("{:?}", double_client.double(context::current(), i).await?);
}
Ok(())
}

View File

@@ -0,0 +1,150 @@
// Copyright 2023 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use futures::prelude::*;
use rustls_pemfile::certs;
use std::io::{BufReader, Cursor};
use std::net::{IpAddr, Ipv4Addr};
use tokio_rustls::rustls::server::AllowAnyAuthenticatedClient;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::net::TcpStream;
use tokio_rustls::rustls::{self, RootCertStore};
use tokio_rustls::{TlsAcceptor, TlsConnector};
use tarpc::context::Context;
use tarpc::serde_transport as transport;
use tarpc::server::{BaseChannel, Channel};
use tarpc::tokio_serde::formats::Bincode;
use tarpc::tokio_util::codec::length_delimited::LengthDelimitedCodec;
#[tarpc::service]
pub trait PingService {
async fn ping() -> String;
}
#[derive(Clone)]
struct Service;
impl PingService for Service {
async fn ping(self, _: Context) -> String {
"🔒".to_owned()
}
}
// certs were generated with openssl 3 https://github.com/rustls/rustls/tree/main/test-ca
// used on client-side for server tls
const END_CHAIN: &str = include_str!("certs/eddsa/end.chain");
// used on client-side for client-auth
const CLIENT_PRIVATEKEY_CLIENT_AUTH: &str = include_str!("certs/eddsa/client.key");
const CLIENT_CERT_CLIENT_AUTH: &str = include_str!("certs/eddsa/client.cert");
// used on server-side for server tls
const END_CERT: &str = include_str!("certs/eddsa/end.cert");
const END_PRIVATEKEY: &str = include_str!("certs/eddsa/end.key");
// used on server-side for client-auth
const CLIENT_CHAIN_CLIENT_AUTH: &str = include_str!("certs/eddsa/client.chain");
pub fn load_certs(data: &str) -> Vec<rustls::Certificate> {
certs(&mut BufReader::new(Cursor::new(data)))
.unwrap()
.into_iter()
.map(rustls::Certificate)
.collect()
}
pub fn load_private_key(key: &str) -> rustls::PrivateKey {
let mut reader = BufReader::new(Cursor::new(key));
loop {
match rustls_pemfile::read_one(&mut reader).expect("cannot parse private key .pem file") {
Some(rustls_pemfile::Item::RSAKey(key)) => return rustls::PrivateKey(key),
Some(rustls_pemfile::Item::PKCS8Key(key)) => return rustls::PrivateKey(key),
Some(rustls_pemfile::Item::ECKey(key)) => return rustls::PrivateKey(key),
None => break,
_ => {}
}
}
panic!("no keys found in {:?} (encrypted keys not supported)", key);
}
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
// -------------------- start here to setup tls tcp tokio stream --------------------------
// ref certs and loading from: https://github.com/tokio-rs/tls/blob/master/tokio-rustls/tests/test.rs
// ref basic tls server setup from: https://github.com/tokio-rs/tls/blob/master/tokio-rustls/examples/server/src/main.rs
let cert = load_certs(END_CERT);
let key = load_private_key(END_PRIVATEKEY);
let server_addr = (IpAddr::V4(Ipv4Addr::LOCALHOST), 5000);
// ------------- server side client_auth cert loading start
let mut client_auth_roots = RootCertStore::empty();
for root in load_certs(CLIENT_CHAIN_CLIENT_AUTH) {
client_auth_roots.add(&root).unwrap();
}
let client_auth = AllowAnyAuthenticatedClient::new(client_auth_roots);
// ------------- server side client_auth cert loading end
let config = rustls::ServerConfig::builder()
.with_safe_defaults()
.with_client_cert_verifier(client_auth) // use .with_no_client_auth() instead if you don't want client-auth
.with_single_cert(cert, key)
.unwrap();
let acceptor = TlsAcceptor::from(Arc::new(config));
let listener = TcpListener::bind(&server_addr).await.unwrap();
let codec_builder = LengthDelimitedCodec::builder();
// ref ./custom_transport.rs server side
tokio::spawn(async move {
loop {
let (stream, _peer_addr) = listener.accept().await.unwrap();
let tls_stream = acceptor.accept(stream).await.unwrap();
let framed = codec_builder.new_framed(tls_stream);
let transport = transport::new(framed, Bincode::default());
let fut = BaseChannel::with_defaults(transport)
.execute(Service.serve())
.for_each(spawn);
tokio::spawn(fut);
}
});
// ---------------------- client connection ---------------------
// tls client connection from https://github.com/tokio-rs/tls/blob/master/tokio-rustls/examples/client/src/main.rs
let mut root_store = rustls::RootCertStore::empty();
for root in load_certs(END_CHAIN) {
root_store.add(&root).unwrap();
}
let client_auth_private_key = load_private_key(CLIENT_PRIVATEKEY_CLIENT_AUTH);
let client_auth_certs = load_certs(CLIENT_CERT_CLIENT_AUTH);
let config = rustls::ClientConfig::builder()
.with_safe_defaults()
.with_root_certificates(root_store)
.with_single_cert(client_auth_certs, client_auth_private_key)?; // use .with_no_client_auth() instead if you don't want client-auth
let domain = rustls::ServerName::try_from("localhost")?;
let connector = TlsConnector::from(Arc::new(config));
let stream = TcpStream::connect(server_addr).await?;
let stream = connector.connect(domain, stream).await?;
let transport = transport::new(codec_builder.new_framed(stream), Bincode::default());
let answer = PingServiceClient::new(Default::default(), transport)
.spawn()
.ping(tarpc::context::current())
.await?;
println!("ping answer: {answer}");
Ok(())
}

190
tarpc/examples/tracing.rs Normal file
View File

@@ -0,0 +1,190 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use crate::{
add::{Add as AddService, AddStub},
double::Double as DoubleService,
};
use futures::{future, prelude::*};
use std::{
io,
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
};
use tarpc::{
client::{
self,
stub::{load_balance, retry},
RpcError,
},
context, serde_transport,
server::{
incoming::{spawn_incoming, Incoming},
request_hook::{self, BeforeRequestList},
BaseChannel,
},
tokio_serde::formats::Json,
ClientMessage, Response, ServerError, Transport,
};
use tokio::net::TcpStream;
use tracing_subscriber::prelude::*;
pub mod add {
#[tarpc::service]
pub trait Add {
/// Add two ints together.
async fn add(x: i32, y: i32) -> i32;
}
}
pub mod double {
#[tarpc::service]
pub trait Double {
/// 2 * x
async fn double(x: i32) -> Result<i32, String>;
}
}
#[derive(Clone)]
struct AddServer;
impl AddService for AddServer {
async fn add(self, _: context::Context, x: i32, y: i32) -> i32 {
x + y
}
}
#[derive(Clone)]
struct DoubleServer<Stub> {
add_client: add::AddClient<Stub>,
}
impl<Stub> DoubleService for DoubleServer<Stub>
where
Stub: AddStub + Clone + Send + Sync + 'static,
{
async fn double(self, _: context::Context, x: i32) -> Result<i32, String> {
self.add_client
.add(context::current(), x, x)
.await
.map_err(|e| e.to_string())
}
}
fn init_tracing(service_name: &str) -> anyhow::Result<()> {
let tracer = opentelemetry_jaeger::new_agent_pipeline()
.with_service_name(service_name)
.with_auto_split_batch(true)
.with_max_packet_size(2usize.pow(13))
.install_batch(opentelemetry::runtime::Tokio)?;
tracing_subscriber::registry()
.with(tracing_subscriber::EnvFilter::from_default_env())
.with(tracing_subscriber::fmt::layer())
.with(tracing_opentelemetry::layer().with_tracer(tracer))
.try_init()?;
Ok(())
}
async fn listen_on_random_port<Item, SinkItem>() -> anyhow::Result<(
impl Stream<Item = serde_transport::Transport<TcpStream, Item, SinkItem, Json<Item, SinkItem>>>,
std::net::SocketAddr,
)>
where
Item: for<'de> serde::Deserialize<'de>,
SinkItem: serde::Serialize,
{
let listener = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()))
.take(1);
let addr = listener.get_ref().get_ref().local_addr();
Ok((listener, addr))
}
fn make_stub<Req, Resp, const N: usize>(
backends: [impl Transport<ClientMessage<Arc<Req>>, Response<Resp>> + Send + Sync + 'static; N],
) -> retry::Retry<
impl Fn(&Result<Resp, RpcError>, u32) -> bool + Clone,
load_balance::RoundRobin<client::Channel<Arc<Req>, Resp>>,
>
where
Req: Send + Sync + 'static,
Resp: Send + Sync + 'static,
{
let stub = load_balance::RoundRobin::new(
backends
.into_iter()
.map(|transport| tarpc::client::new(client::Config::default(), transport).spawn())
.collect(),
);
let stub = retry::Retry::new(stub, |resp, attempts| {
if let Err(e) = resp {
tracing::warn!("Got an error: {e:?}");
attempts < 3
} else {
false
}
});
stub
}
#[tokio::main]
async fn main() -> anyhow::Result<()> {
init_tracing("tarpc_tracing_example")?;
let (add_listener1, addr1) = listen_on_random_port().await?;
let (add_listener2, addr2) = listen_on_random_port().await?;
let something_bad_happened = Arc::new(AtomicBool::new(false));
let server = request_hook::before()
.then_fn(move |_: &mut _, _: &_| {
let something_bad_happened = something_bad_happened.clone();
async move {
if something_bad_happened.fetch_xor(true, Ordering::Relaxed) {
Err(ServerError::new(
io::ErrorKind::NotFound,
"Gamma Ray!".into(),
))
} else {
Ok(())
}
}
})
.serving(AddServer.serve());
let add_server = add_listener1
.chain(add_listener2)
.map(BaseChannel::with_defaults);
tokio::spawn(spawn_incoming(add_server.execute(server)));
let add_client = add::AddClient::from(make_stub([
tarpc::serde_transport::tcp::connect(addr1, Json::default).await?,
tarpc::serde_transport::tcp::connect(addr2, Json::default).await?,
]));
let double_listener = tarpc::serde_transport::tcp::listen("localhost:0", Json::default)
.await?
.filter_map(|r| future::ready(r.ok()));
let addr = double_listener.get_ref().local_addr();
let double_server = double_listener.map(BaseChannel::with_defaults).take(1);
let server = DoubleServer { add_client }.serve();
tokio::spawn(spawn_incoming(double_server.execute(server)));
let to_double_server = tarpc::serde_transport::tcp::connect(addr, Json::default).await?;
let double_client =
double::DoubleClient::new(client::Config::default(), to_double_server).spawn();
let ctx = context::current();
for _ in 1..=5 {
tracing::info!("{:?}", double_client.double(ctx, 1).await?);
}
opentelemetry::global::shutdown_tracer_provider();
Ok(())
}

View File

@@ -0,0 +1,49 @@
use futures::{prelude::*, task::*};
use std::pin::Pin;
use tokio::sync::mpsc;
/// Sends request cancellation signals.
#[derive(Debug, Clone)]
pub struct RequestCancellation(mpsc::UnboundedSender<u64>);
/// A stream of IDs of requests that have been canceled.
#[derive(Debug)]
pub struct CanceledRequests(mpsc::UnboundedReceiver<u64>);
/// Returns a channel to send request cancellation messages.
pub fn cancellations() -> (RequestCancellation, CanceledRequests) {
// Unbounded because messages are sent in the drop fn. This is fine, because it's still
// bounded by the number of in-flight requests.
let (tx, rx) = mpsc::unbounded_channel();
(RequestCancellation(tx), CanceledRequests(rx))
}
impl RequestCancellation {
/// Cancels the request with ID `request_id`.
///
/// No validation is done of `request_id`. There is no way to know if the request id provided
/// corresponds to a request actually tracked by the backing channel. `RequestCancellation` is
/// a one-way communication channel.
///
/// Once request data is cleaned up, a response will never be received by the client. This is
/// useful primarily when request processing ends prematurely for requests with long deadlines
/// which would otherwise continue to be tracked by the backing channel—a kind of leak.
pub fn cancel(&self, request_id: u64) {
let _ = self.0.send(request_id);
}
}
impl CanceledRequests {
/// Polls for a cancelled request.
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<u64>> {
self.0.poll_recv(cx)
}
}
impl Stream for CanceledRequests {
type Item = u64;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<u64>> {
self.poll_recv(cx)
}
}

1065
tarpc/src/client.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,137 @@
use crate::{
context,
util::{Compact, TimeUntil},
};
use fnv::FnvHashMap;
use std::{
collections::hash_map,
task::{Context, Poll},
};
use tokio::sync::oneshot;
use tokio_util::time::delay_queue::{self, DelayQueue};
use tracing::Span;
/// Requests already written to the wire that haven't yet received responses.
#[derive(Debug)]
pub struct InFlightRequests<Resp> {
request_data: FnvHashMap<u64, RequestData<Resp>>,
deadlines: DelayQueue<u64>,
}
impl<Resp> Default for InFlightRequests<Resp> {
fn default() -> Self {
Self {
request_data: Default::default(),
deadlines: Default::default(),
}
}
}
#[derive(Debug)]
struct RequestData<Res> {
ctx: context::Context,
span: Span,
response_completion: oneshot::Sender<Res>,
/// The key to remove the timer for the request's deadline.
deadline_key: delay_queue::Key,
}
/// An error returned when an attempt is made to insert a request with an ID that is already in
/// use.
#[derive(Debug)]
pub struct AlreadyExistsError;
impl<Res> InFlightRequests<Res> {
/// Returns the number of in-flight requests.
pub fn len(&self) -> usize {
self.request_data.len()
}
/// Returns true iff there are no requests in flight.
pub fn is_empty(&self) -> bool {
self.request_data.is_empty()
}
/// Starts a request, unless a request with the same ID is already in flight.
pub fn insert_request(
&mut self,
request_id: u64,
ctx: context::Context,
span: Span,
response_completion: oneshot::Sender<Res>,
) -> Result<(), AlreadyExistsError> {
match self.request_data.entry(request_id) {
hash_map::Entry::Vacant(vacant) => {
let timeout = ctx.deadline.time_until();
let deadline_key = self.deadlines.insert(request_id, timeout);
vacant.insert(RequestData {
ctx,
span,
response_completion,
deadline_key,
});
Ok(())
}
hash_map::Entry::Occupied(_) => Err(AlreadyExistsError),
}
}
/// Removes a request without aborting. Returns true iff the request was found.
pub fn complete_request(&mut self, request_id: u64, result: Res) -> Option<Span> {
if let Some(request_data) = self.request_data.remove(&request_id) {
self.request_data.compact(0.1);
self.deadlines.remove(&request_data.deadline_key);
let _ = request_data.response_completion.send(result);
return Some(request_data.span);
}
tracing::debug!("No in-flight request found for request_id = {request_id}.");
// If the response completion was absent, then the request was already canceled.
None
}
/// Completes all requests using the provided function.
/// Returns Spans for all completes requests.
pub fn complete_all_requests<'a>(
&'a mut self,
mut result: impl FnMut() -> Res + 'a,
) -> impl Iterator<Item = Span> + 'a {
self.deadlines.clear();
self.request_data.drain().map(move |(_, request_data)| {
let _ = request_data.response_completion.send(result());
request_data.span
})
}
/// Cancels a request without completing (typically used when a request handle was dropped
/// before the request completed).
pub fn cancel_request(&mut self, request_id: u64) -> Option<(context::Context, Span)> {
if let Some(request_data) = self.request_data.remove(&request_id) {
self.request_data.compact(0.1);
self.deadlines.remove(&request_data.deadline_key);
Some((request_data.ctx, request_data.span))
} else {
None
}
}
/// Yields a request that has expired, completing it with a TimedOut error.
/// The caller should send cancellation messages for any yielded request ID.
pub fn poll_expired(
&mut self,
cx: &mut Context,
expired_error: impl Fn() -> Res,
) -> Poll<Option<u64>> {
self.deadlines.poll_expired(cx).map(|expired| {
let request_id = expired?.into_inner();
if let Some(request_data) = self.request_data.remove(&request_id) {
let _entered = request_data.span.enter();
tracing::error!("DeadlineExceeded");
self.request_data.compact(0.1);
let _ = request_data.response_completion.send(expired_error());
}
Some(request_id)
})
}
}

45
tarpc/src/client/stub.rs Normal file
View File

@@ -0,0 +1,45 @@
//! Provides a Stub trait, implemented by types that can call remote services.
use crate::{
client::{Channel, RpcError},
context,
};
pub mod load_balance;
pub mod retry;
#[cfg(test)]
mod mock;
/// A connection to a remote service.
/// Calls the service with requests of type `Req` and receives responses of type `Resp`.
#[allow(async_fn_in_trait)]
pub trait Stub {
/// The service request type.
type Req;
/// The service response type.
type Resp;
/// Calls a remote service.
async fn call(
&self,
ctx: context::Context,
request_name: &'static str,
request: Self::Req,
) -> Result<Self::Resp, RpcError>;
}
impl<Req, Resp> Stub for Channel<Req, Resp> {
type Req = Req;
type Resp = Resp;
async fn call(
&self,
ctx: context::Context,
request_name: &'static str,
request: Req,
) -> Result<Self::Resp, RpcError> {
Self::call(self, ctx, request_name, request).await
}
}

View File

@@ -0,0 +1,279 @@
//! Provides load-balancing [Stubs](crate::client::stub::Stub).
pub use consistent_hash::ConsistentHash;
pub use round_robin::RoundRobin;
/// Provides a stub that load-balances with a simple round-robin strategy.
mod round_robin {
use crate::{
client::{stub, RpcError},
context,
};
use cycle::AtomicCycle;
impl<Stub> stub::Stub for RoundRobin<Stub>
where
Stub: stub::Stub,
{
type Req = Stub::Req;
type Resp = Stub::Resp;
async fn call(
&self,
ctx: context::Context,
request_name: &'static str,
request: Self::Req,
) -> Result<Stub::Resp, RpcError> {
let next = self.stubs.next();
next.call(ctx, request_name, request).await
}
}
/// A Stub that load-balances across backing stubs by round robin.
#[derive(Clone, Debug)]
pub struct RoundRobin<Stub> {
stubs: AtomicCycle<Stub>,
}
impl<Stub> RoundRobin<Stub>
where
Stub: stub::Stub,
{
/// Returns a new RoundRobin stub.
pub fn new(stubs: Vec<Stub>) -> Self {
Self {
stubs: AtomicCycle::new(stubs),
}
}
}
mod cycle {
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
/// Cycles endlessly and atomically over a collection of elements of type T.
#[derive(Clone, Debug)]
pub struct AtomicCycle<T>(Arc<State<T>>);
#[derive(Debug)]
struct State<T> {
elements: Vec<T>,
next: AtomicUsize,
}
impl<T> AtomicCycle<T> {
pub fn new(elements: Vec<T>) -> Self {
Self(Arc::new(State {
elements,
next: Default::default(),
}))
}
pub fn next(&self) -> &T {
self.0.next()
}
}
impl<T> State<T> {
pub fn next(&self) -> &T {
let next = self.next.fetch_add(1, Ordering::Relaxed);
&self.elements[next % self.elements.len()]
}
}
#[test]
fn test_cycle() {
let cycle = AtomicCycle::new(vec![1, 2, 3]);
assert_eq!(cycle.next(), &1);
assert_eq!(cycle.next(), &2);
assert_eq!(cycle.next(), &3);
assert_eq!(cycle.next(), &1);
}
}
}
/// Provides a stub that load-balances with a consistent hashing strategy.
///
/// Each request is hashed, then mapped to a stub based on the hash. Equivalent requests will use
/// the same stub.
mod consistent_hash {
use crate::{
client::{stub, RpcError},
context,
};
use std::{
collections::hash_map::RandomState,
hash::{BuildHasher, Hash, Hasher},
num::TryFromIntError,
};
impl<Stub, S> stub::Stub for ConsistentHash<Stub, S>
where
Stub: stub::Stub,
Stub::Req: Hash,
S: BuildHasher,
{
type Req = Stub::Req;
type Resp = Stub::Resp;
async fn call(
&self,
ctx: context::Context,
request_name: &'static str,
request: Self::Req,
) -> Result<Stub::Resp, RpcError> {
let index = usize::try_from(self.hash_request(&request) % self.stubs_len).expect(
"invariant broken: stubs_len is not larger than a usize, \
so the hash modulo stubs_len should always fit in a usize",
);
let next = &self.stubs[index];
next.call(ctx, request_name, request).await
}
}
/// A Stub that load-balances across backing stubs by round robin.
#[derive(Clone, Debug)]
pub struct ConsistentHash<Stub, S = RandomState> {
stubs: Vec<Stub>,
stubs_len: u64,
hasher: S,
}
impl<Stub> ConsistentHash<Stub, RandomState>
where
Stub: stub::Stub,
Stub::Req: Hash,
{
/// Returns a new RoundRobin stub.
/// Returns an err if the length of `stubs` overflows a u64.
pub fn new(stubs: Vec<Stub>) -> Result<Self, TryFromIntError> {
Ok(Self {
stubs_len: stubs.len().try_into()?,
stubs,
hasher: RandomState::new(),
})
}
}
impl<Stub, S> ConsistentHash<Stub, S>
where
Stub: stub::Stub,
Stub::Req: Hash,
S: BuildHasher,
{
/// Returns a new RoundRobin stub.
/// Returns an err if the length of `stubs` overflows a u64.
pub fn with_hasher(stubs: Vec<Stub>, hasher: S) -> Result<Self, TryFromIntError> {
Ok(Self {
stubs_len: stubs.len().try_into()?,
stubs,
hasher,
})
}
fn hash_request(&self, req: &Stub::Req) -> u64 {
let mut hasher = self.hasher.build_hasher();
req.hash(&mut hasher);
hasher.finish()
}
}
#[cfg(test)]
mod tests {
use super::ConsistentHash;
use crate::{
client::stub::{mock::Mock, Stub},
context,
};
use std::{
collections::HashMap,
hash::{BuildHasher, Hash, Hasher},
rc::Rc,
};
#[tokio::test]
async fn test() -> anyhow::Result<()> {
let stub = ConsistentHash::<_, FakeHasherBuilder>::with_hasher(
vec![
// For easier reading of the assertions made in this test, each Mock's response
// value is equal to a hash value that should map to its index: 3 % 3 = 0, 1 %
// 3 = 1, etc.
Mock::new([('a', 3), ('b', 3), ('c', 3)]),
Mock::new([('a', 1), ('b', 1), ('c', 1)]),
Mock::new([('a', 2), ('b', 2), ('c', 2)]),
],
FakeHasherBuilder::new([('a', 1), ('b', 2), ('c', 3)]),
)?;
for _ in 0..2 {
let resp = stub.call(context::current(), "", 'a').await?;
assert_eq!(resp, 1);
let resp = stub.call(context::current(), "", 'b').await?;
assert_eq!(resp, 2);
let resp = stub.call(context::current(), "", 'c').await?;
assert_eq!(resp, 3);
}
Ok(())
}
struct HashRecorder(Vec<u8>);
impl Hasher for HashRecorder {
fn write(&mut self, bytes: &[u8]) {
self.0 = Vec::from(bytes);
}
fn finish(&self) -> u64 {
0
}
}
struct FakeHasherBuilder {
recorded_hashes: Rc<HashMap<Vec<u8>, u64>>,
}
struct FakeHasher {
recorded_hashes: Rc<HashMap<Vec<u8>, u64>>,
output: u64,
}
impl BuildHasher for FakeHasherBuilder {
type Hasher = FakeHasher;
fn build_hasher(&self) -> Self::Hasher {
FakeHasher {
recorded_hashes: self.recorded_hashes.clone(),
output: 0,
}
}
}
impl FakeHasherBuilder {
fn new<T: Hash, const N: usize>(fake_hashes: [(T, u64); N]) -> Self {
let mut recorded_hashes = HashMap::new();
for (to_hash, fake_hash) in fake_hashes {
let mut recorder = HashRecorder(vec![]);
to_hash.hash(&mut recorder);
recorded_hashes.insert(recorder.0, fake_hash);
}
Self {
recorded_hashes: Rc::new(recorded_hashes),
}
}
}
impl Hasher for FakeHasher {
fn write(&mut self, bytes: &[u8]) {
if let Some(hash) = self.recorded_hashes.get(bytes) {
self.output = *hash;
}
}
fn finish(&self) -> u64 {
self.output
}
}
}
}

View File

@@ -0,0 +1,49 @@
use crate::{
client::{stub::Stub, RpcError},
context, ServerError,
};
use std::{collections::HashMap, hash::Hash, io};
/// A mock stub that returns user-specified responses.
pub struct Mock<Req, Resp> {
responses: HashMap<Req, Resp>,
}
impl<Req, Resp> Mock<Req, Resp>
where
Req: Eq + Hash,
{
/// Returns a new mock, mocking the specified (request, response) pairs.
pub fn new<const N: usize>(responses: [(Req, Resp); N]) -> Self {
Self {
responses: HashMap::from(responses),
}
}
}
impl<Req, Resp> Stub for Mock<Req, Resp>
where
Req: Eq + Hash,
Resp: Clone,
{
type Req = Req;
type Resp = Resp;
async fn call(
&self,
_: context::Context,
_: &'static str,
request: Self::Req,
) -> Result<Resp, RpcError> {
self.responses
.get(&request)
.cloned()
.map(Ok)
.unwrap_or_else(|| {
Err(RpcError::Server(ServerError {
kind: io::ErrorKind::NotFound,
detail: "mock (request, response) entry not found".into(),
}))
})
}
}

View File

@@ -0,0 +1,56 @@
//! Provides a stub that retries requests based on response contents..
use crate::{
client::{stub, RpcError},
context,
};
use std::sync::Arc;
impl<Stub, Req, F> stub::Stub for Retry<F, Stub>
where
Stub: stub::Stub<Req = Arc<Req>>,
F: Fn(&Result<Stub::Resp, RpcError>, u32) -> bool,
{
type Req = Req;
type Resp = Stub::Resp;
async fn call(
&self,
ctx: context::Context,
request_name: &'static str,
request: Self::Req,
) -> Result<Stub::Resp, RpcError> {
let request = Arc::new(request);
for i in 1.. {
let result = self
.stub
.call(ctx, request_name, Arc::clone(&request))
.await;
if (self.should_retry)(&result, i) {
tracing::trace!("Retrying on attempt {i}");
continue;
}
return result;
}
unreachable!("Wow, that was a lot of attempts!");
}
}
/// A Stub that retries requests based on response contents.
/// Note: to use this stub with Serde serialization, the "rc" feature of Serde needs to be enabled.
#[derive(Clone, Debug)]
pub struct Retry<F, Stub> {
should_retry: F,
stub: Stub,
}
impl<Stub, Req, F> Retry<F, Stub>
where
Stub: stub::Stub<Req = Arc<Req>>,
F: Fn(&Result<Stub::Resp, RpcError>, u32) -> bool,
{
/// Creates a new Retry stub that delegates calls to the underlying `stub`.
pub fn new(stub: Stub, should_retry: F) -> Self {
Self { stub, should_retry }
}
}

152
tarpc/src/context.rs Normal file
View File

@@ -0,0 +1,152 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a request context that carries a deadline and trace context. This context is sent from
//! client to server and is used by the server to enforce response deadlines.
use crate::trace::{self, TraceId};
use opentelemetry::trace::TraceContextExt;
use static_assertions::assert_impl_all;
use std::{
convert::TryFrom,
time::{Duration, SystemTime},
};
use tracing_opentelemetry::OpenTelemetrySpanExt;
/// A request context that carries request-scoped information like deadlines and trace information.
/// It is sent from client to server and is used by the server to enforce response deadlines.
///
/// The context should not be stored directly in a server implementation, because the context will
/// be different for each request in scope.
#[derive(Clone, Copy, Debug)]
#[non_exhaustive]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Context {
/// When the client expects the request to be complete by. The server should cancel the request
/// if it is not complete by this time.
#[cfg_attr(feature = "serde1", serde(default = "ten_seconds_from_now"))]
// Serialized as a Duration to prevent clock skew issues.
#[cfg_attr(feature = "serde1", serde(with = "absolute_to_relative_time"))]
pub deadline: SystemTime,
/// Uniquely identifies requests originating from the same source.
/// When a service handles a request by making requests itself, those requests should
/// include the same `trace_id` as that included on the original request. This way,
/// users can trace related actions across a distributed system.
pub trace_context: trace::Context,
}
#[cfg(feature = "serde1")]
mod absolute_to_relative_time {
pub use serde::{Deserialize, Deserializer, Serialize, Serializer};
pub use std::time::{Duration, SystemTime};
pub fn serialize<S>(deadline: &SystemTime, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let deadline = deadline
.duration_since(SystemTime::now())
.unwrap_or(Duration::ZERO);
deadline.serialize(serializer)
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<SystemTime, D::Error>
where
D: Deserializer<'de>,
{
let deadline = Duration::deserialize(deserializer)?;
Ok(SystemTime::now() + deadline)
}
#[cfg(test)]
#[derive(serde::Serialize, serde::Deserialize)]
struct AbsoluteToRelative(#[serde(with = "self")] SystemTime);
#[test]
fn test_serialize() {
let now = SystemTime::now();
let deadline = now + Duration::from_secs(10);
let serialized_deadline = bincode::serialize(&AbsoluteToRelative(deadline)).unwrap();
let deserialized_deadline: Duration = bincode::deserialize(&serialized_deadline).unwrap();
// TODO: how to avoid flakiness?
assert!(deserialized_deadline > Duration::from_secs(9));
}
#[test]
fn test_deserialize() {
let deadline = Duration::from_secs(10);
let serialized_deadline = bincode::serialize(&deadline).unwrap();
let AbsoluteToRelative(deserialized_deadline) =
bincode::deserialize(&serialized_deadline).unwrap();
// TODO: how to avoid flakiness?
assert!(deserialized_deadline > SystemTime::now() + Duration::from_secs(9));
}
}
assert_impl_all!(Context: Send, Sync);
fn ten_seconds_from_now() -> SystemTime {
SystemTime::now() + Duration::from_secs(10)
}
/// Returns the context for the current request, or a default Context if no request is active.
pub fn current() -> Context {
Context::current()
}
#[derive(Clone)]
struct Deadline(SystemTime);
impl Default for Deadline {
fn default() -> Self {
Self(ten_seconds_from_now())
}
}
impl Context {
/// Returns the context for the current request, or a default Context if no request is active.
pub fn current() -> Self {
let span = tracing::Span::current();
Self {
trace_context: trace::Context::try_from(&span)
.unwrap_or_else(|_| trace::Context::default()),
deadline: span
.context()
.get::<Deadline>()
.cloned()
.unwrap_or_default()
.0,
}
}
/// Returns the ID of the request-scoped trace.
pub fn trace_id(&self) -> &TraceId {
&self.trace_context.trace_id
}
}
/// An extension trait for [`tracing::Span`] for propagating tarpc Contexts.
pub(crate) trait SpanExt {
/// Sets the given context on this span. Newly-created spans will be children of the given
/// context's trace context.
fn set_context(&self, context: &Context);
}
impl SpanExt for tracing::Span {
fn set_context(&self, context: &Context) {
self.set_parent(
opentelemetry::Context::new()
.with_remote_span_context(opentelemetry::trace::SpanContext::new(
opentelemetry::trace::TraceId::from(context.trace_context.trace_id),
opentelemetry::trace::SpanId::from(context.trace_context.span_id),
opentelemetry::trace::TraceFlags::from(context.trace_context.sampling_decision),
true,
opentelemetry::trace::TraceState::default(),
))
.with_value(Deadline(context.deadline)),
);
}
}

View File

@@ -3,10 +3,6 @@
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! [![Latest Version](https://img.shields.io/crates/v/tarpc.svg)](https://crates.io/crates/tarpc)
//! [![Join the chat at https://gitter.im/tarpc/Lobby](https://badges.gitter.im/tarpc/Lobby.svg)](https://gitter.im/tarpc/Lobby?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge)
//!
//! *Disclaimer*: This is not an official Google product.
//!
//! tarpc is an RPC framework for rust with a focus on ease of use. Defining a
@@ -31,7 +27,7 @@
//! process, and no context switching between different languages.
//!
//! Some other features of tarpc:
//! - Pluggable transport: any type impling `Stream<Item = Request> + Sink<Response>` can be
//! - Pluggable transport: any type implementing `Stream<Item = Request> + Sink<Response>` can be
//! used as a transport to connect the client and server.
//! - `Send + 'static` optional: if the transport doesn't require it, neither does tarpc!
//! - Cascading cancellation: dropping a request will send a cancellation message to the server.
@@ -42,6 +38,14 @@
//! requests sent by the server that use the request context will propagate the request deadline.
//! For example, if a server is handling a request with a 10s deadline, does 2s of work, then
//! sends a request to another server, that server will see an 8s deadline.
//! - Distributed tracing: tarpc is instrumented with
//! [tracing](https://github.com/tokio-rs/tracing) primitives extended with
//! [OpenTelemetry](https://opentelemetry.io/) traces. Using a compatible tracing subscriber like
//! [Jaeger](https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-jaeger),
//! each RPC can be traced through the client, server, and other dependencies downstream of the
//! server. Even for applications not connected to a distributed tracing collector, the
//! instrumentation can also be ingested by regular loggers like
//! [env_logger](https://github.com/env-logger-rs/env_logger/).
//! - Serde serialization: enabling the `serde1` Cargo feature will make service requests and
//! responses `Serialize + Deserialize`. It's entirely optional, though: in-memory transports can
//! be used, as well, so the price of serialization doesn't have to be paid when it's not needed.
@@ -50,7 +54,7 @@
//! Add to your `Cargo.toml` dependencies:
//!
//! ```toml
//! tarpc = "0.20.0"
//! tarpc = "0.29"
//! ```
//!
//! The `tarpc::service` attribute expands to a collection of items that form an rpc service.
@@ -59,12 +63,14 @@
//!
//! ## Example
//!
//! For this example, in addition to tarpc, also add two other dependencies to
//! This example uses [tokio](https://tokio.rs), so add the following dependencies to
//! your `Cargo.toml`:
//!
//! ```toml
//! anyhow = "1.0"
//! futures = "0.3"
//! tokio = "0.2"
//! tarpc = { version = "0.29", features = ["tokio1"] }
//! tokio = { version = "1.0", features = ["macros"] }
//! ```
//!
//! In the following example, we use an in-process channel for communication between
@@ -82,9 +88,8 @@
//! };
//! use tarpc::{
//! client, context,
//! server::{self, Handler},
//! server::{self, incoming::Incoming, Channel},
//! };
//! use std::io;
//!
//! // This is the service definition. It looks a lot like a trait definition.
//! // It defines one RPC, hello, which takes one arg, name, and returns a String.
@@ -106,9 +111,8 @@
//! # };
//! # use tarpc::{
//! # client, context,
//! # server::{self, Handler},
//! # server::{self, incoming::Incoming},
//! # };
//! # use std::io;
//! # // This is the service definition. It looks a lot like a trait definition.
//! # // It defines one RPC, hello, which takes one arg, name, and returns a String.
//! # #[tarpc::service]
@@ -122,19 +126,15 @@
//! struct HelloServer;
//!
//! impl World for HelloServer {
//! // Each defined rpc generates two items in the trait, a fn that serves the RPC, and
//! // an associated type representing the future output by the fn.
//!
//! type HelloFut = Ready<String>;
//!
//! fn hello(self, _: context::Context, name: String) -> Self::HelloFut {
//! future::ready(format!("Hello, {}!", name))
//! // Each defined rpc generates an async fn that serves the RPC
//! async fn hello(self, _: context::Context, name: String) -> String {
//! format!("Hello, {name}!")
//! }
//! }
//! ```
//!
//! Lastly let's write our `main` that will start the server. While this example uses an
//! [in-process channel](rpc::transport::channel), tarpc also ships a generic [`serde_transport`]
//! [in-process channel](transport::channel), tarpc also ships a generic [`serde_transport`]
//! behind the `serde-transport` feature, with additional [TCP](serde_transport::tcp) functionality
//! available behind the `tcp` feature.
//!
@@ -146,9 +146,8 @@
//! # };
//! # use tarpc::{
//! # client, context,
//! # server::{self, Handler},
//! # server::{self, Channel},
//! # };
//! # use std::io;
//! # // This is the service definition. It looks a lot like a trait definition.
//! # // It defines one RPC, hello, which takes one arg, name, and returns a String.
//! # #[tarpc::service]
@@ -161,35 +160,36 @@
//! # #[derive(Clone)]
//! # struct HelloServer;
//! # impl World for HelloServer {
//! # // Each defined rpc generates two items in the trait, a fn that serves the RPC, and
//! # // an associated type representing the future output by the fn.
//! # type HelloFut = Ready<String>;
//! # fn hello(self, _: context::Context, name: String) -> Self::HelloFut {
//! # future::ready(format!("Hello, {}!", name))
//! // Each defined rpc generates an async fn that serves the RPC
//! # async fn hello(self, _: context::Context, name: String) -> String {
//! # format!("Hello, {name}!")
//! # }
//! # }
//! # #[cfg(not(feature = "tokio1"))]
//! # fn main() {}
//! # #[cfg(feature = "tokio1")]
//! #[tokio::main]
//! async fn main() -> io::Result<()> {
//! async fn main() -> anyhow::Result<()> {
//! let (client_transport, server_transport) = tarpc::transport::channel::unbounded();
//!
//! let server = server::new(server::Config::default())
//! // incoming() takes a stream of transports such as would be returned by
//! // TcpListener::incoming (but a stream instead of an iterator).
//! .incoming(stream::once(future::ready(server_transport)))
//! .respond_with(HelloServer.serve());
//! let server = server::BaseChannel::with_defaults(server_transport);
//! tokio::spawn(
//! server.execute(HelloServer.serve())
//! // Handle all requests concurrently.
//! .for_each(|response| async move {
//! tokio::spawn(response);
//! }));
//!
//! tokio::spawn(server);
//!
//! // WorldClient is generated by the macro. It has a constructor `new` that takes a config and
//! // any Transport as input
//! let mut client = WorldClient::new(client::Config::default(), client_transport).spawn()?;
//! // WorldClient is generated by the #[tarpc::service] attribute. It has a constructor `new`
//! // that takes a config and any Transport as input.
//! let mut client = WorldClient::new(client::Config::default(), client_transport).spawn();
//!
//! // The client has an RPC method for each RPC defined in the annotated trait. It takes the same
//! // args as defined, with the addition of a Context, which is always the first arg. The Context
//! // specifies a deadline and trace information which can be helpful in debugging requests.
//! let hello = client.hello(context::current(), "Stim".to_string()).await?;
//!
//! println!("{}", hello);
//! println!("{hello}");
//!
//! Ok(())
//! }
@@ -199,23 +199,32 @@
//!
//! Use `cargo doc` as you normally would to see the documentation created for all
//! items expanded by a `service!` invocation.
#![deny(missing_docs)]
#![allow(clippy::type_complexity)]
#![cfg_attr(docsrs, feature(doc_cfg))]
pub mod rpc;
pub use rpc::*;
#[cfg(feature = "serde1")]
#[doc(hidden)]
pub use serde;
#[cfg(feature = "serde-transport")]
pub use {tokio_serde, tokio_util};
#[cfg(feature = "serde-transport")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde-transport")))]
pub mod serde_transport;
pub mod trace;
#[cfg(feature = "serde1")]
pub use tarpc_plugins::derive_serde;
/// The main macro that creates RPC services.
///
/// Rpc methods are specified, mirroring trait syntax:
///
/// ```
/// # fn main() {}
/// #[tarpc::service]
/// trait Service {
/// /// Say hello
@@ -234,3 +243,155 @@ pub mod trace;
/// * `Client` -- a client stub with a fn for each RPC.
/// * `fn new_stub` -- creates a new Client stub.
pub use tarpc_plugins::service;
pub(crate) mod cancellations;
pub mod client;
pub mod context;
pub mod server;
pub mod transport;
pub(crate) mod util;
pub use crate::transport::sealed::Transport;
use anyhow::Context as _;
use futures::task::*;
use std::sync::Arc;
use std::{error::Error, fmt::Display, io, time::SystemTime};
/// A message from a client to a server.
#[derive(Debug)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
#[non_exhaustive]
pub enum ClientMessage<T> {
/// A request initiated by a user. The server responds to a request by invoking a
/// service-provided request handler. The handler completes with a [`response`](Response), which
/// the server sends back to the client.
Request(Request<T>),
/// A command to cancel an in-flight request, automatically sent by the client when a response
/// future is dropped.
///
/// When received, the server will immediately cancel the main task (top-level future) of the
/// request handler for the associated request. Any tasks spawned by the request handler will
/// not be canceled, because the framework layer does not
/// know about them.
Cancel {
/// The trace context associates the message with a specific chain of causally-related actions,
/// possibly orchestrated across many distributed systems.
#[cfg_attr(feature = "serde1", serde(default))]
trace_context: trace::Context,
/// The ID of the request to cancel.
request_id: u64,
},
}
/// A request from a client to a server.
#[derive(Clone, Copy, Debug)]
#[non_exhaustive]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Request<T> {
/// Trace context, deadline, and other cross-cutting concerns.
pub context: context::Context,
/// Uniquely identifies the request across all requests sent over a single channel.
pub id: u64,
/// The request body.
pub message: T,
}
/// A response from a server to a client.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
#[non_exhaustive]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Response<T> {
/// The ID of the request being responded to.
pub request_id: u64,
/// The response body, or an error if the request failed.
pub message: Result<T, ServerError>,
}
/// An error indicating the server aborted the request early, e.g., due to request throttling.
#[derive(thiserror::Error, Clone, Debug, PartialEq, Eq, Hash)]
#[error("{kind:?}: {detail}")]
#[non_exhaustive]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct ServerError {
#[cfg_attr(
feature = "serde1",
serde(serialize_with = "util::serde::serialize_io_error_kind_as_u32")
)]
#[cfg_attr(
feature = "serde1",
serde(deserialize_with = "util::serde::deserialize_io_error_kind_from_u32")
)]
/// The type of error that occurred to fail the request.
pub kind: io::ErrorKind,
/// A message describing more detail about the error that occurred.
pub detail: String,
}
/// Critical errors that result in a Channel disconnecting.
#[derive(thiserror::Error, Debug, PartialEq, Eq)]
pub enum ChannelError<E>
where
E: Error + Send + Sync + 'static,
{
/// Could not read from the transport.
#[error("could not read from the transport")]
Read(#[source] Arc<E>),
/// Could not ready the transport for writes.
#[error("could not ready the transport for writes")]
Ready(#[source] E),
/// Could not write to the transport.
#[error("could not write to the transport")]
Write(#[source] E),
/// Could not flush the transport.
#[error("could not flush the transport")]
Flush(#[source] E),
/// Could not close the write end of the transport.
#[error("could not close the write end of the transport")]
Close(#[source] E),
}
impl ServerError {
/// Returns a new server error with `kind` and `detail`.
pub fn new(kind: io::ErrorKind, detail: String) -> ServerError {
Self { kind, detail }
}
}
impl<T> Request<T> {
/// Returns the deadline for this request.
pub fn deadline(&self) -> &SystemTime {
&self.context.deadline
}
}
pub(crate) trait PollContext<T> {
fn context<C>(self, context: C) -> Poll<Option<anyhow::Result<T>>>
where
C: Display + Send + Sync + 'static;
fn with_context<C, F>(self, f: F) -> Poll<Option<anyhow::Result<T>>>
where
C: Display + Send + Sync + 'static,
F: FnOnce() -> C;
}
impl<T, E> PollContext<T> for Poll<Option<Result<T, E>>>
where
E: Error + Send + Sync + 'static,
{
fn context<C>(self, context: C) -> Poll<Option<anyhow::Result<T>>>
where
C: Display + Send + Sync + 'static,
{
self.map(|o| o.map(|r| r.context(context)))
}
fn with_context<C, F>(self, f: F) -> Poll<Option<anyhow::Result<T>>>
where
C: Display + Send + Sync + 'static,
F: FnOnce() -> C,
{
self.map(|o| o.map(|r| r.with_context(f)))
}
}

View File

@@ -1,906 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use crate::{
context,
trace::SpanId,
util::{Compact, TimeUntil},
ClientMessage, PollIo, Request, Response, Transport,
};
use fnv::FnvHashMap;
use futures::{
channel::{mpsc, oneshot},
prelude::*,
ready,
stream::Fuse,
task::*,
};
use log::{debug, info, trace};
use pin_project::{pin_project, pinned_drop};
use std::{
io,
pin::Pin,
sync::{
atomic::{AtomicU64, Ordering},
Arc,
},
};
use super::{Config, NewClient};
/// Handles communication from the client to request dispatch.
#[derive(Debug)]
pub struct Channel<Req, Resp> {
to_dispatch: mpsc::Sender<DispatchRequest<Req, Resp>>,
/// Channel to send a cancel message to the dispatcher.
cancellation: RequestCancellation,
/// The ID to use for the next request to stage.
next_request_id: Arc<AtomicU64>,
}
impl<Req, Resp> Clone for Channel<Req, Resp> {
fn clone(&self) -> Self {
Self {
to_dispatch: self.to_dispatch.clone(),
cancellation: self.cancellation.clone(),
next_request_id: self.next_request_id.clone(),
}
}
}
/// A future returned by [`Channel::send`] that resolves to a server response.
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
struct Send<'a, Req, Resp> {
#[pin]
fut: MapOkDispatchResponse<SendMapErrConnectionReset<'a, Req, Resp>, Resp>,
}
type SendMapErrConnectionReset<'a, Req, Resp> = MapErrConnectionReset<
futures::sink::Send<'a, mpsc::Sender<DispatchRequest<Req, Resp>>, DispatchRequest<Req, Resp>>,
>;
impl<'a, Req, Resp> Future for Send<'a, Req, Resp> {
type Output = io::Result<DispatchResponse<Resp>>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.as_mut().project().fut.poll(cx)
}
}
/// A future returned by [`Channel::call`] that resolves to a server response.
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Call<'a, Req, Resp> {
#[pin]
fut: AndThenIdent<Send<'a, Req, Resp>, DispatchResponse<Resp>>,
}
impl<'a, Req, Resp> Future for Call<'a, Req, Resp> {
type Output = io::Result<Resp>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.as_mut().project().fut.poll(cx)
}
}
impl<Req, Resp> Channel<Req, Resp> {
/// Sends a request to the dispatch task to forward to the server, returning a [`Future`] that
/// resolves when the request is sent (not when the response is received).
fn send(&mut self, mut ctx: context::Context, request: Req) -> Send<Req, Resp> {
// Convert the context to the call context.
ctx.trace_context.parent_id = Some(ctx.trace_context.span_id);
ctx.trace_context.span_id = SpanId::random(&mut rand::thread_rng());
let timeout = ctx.deadline.time_until();
trace!(
"[{}] Queuing request with timeout {:?}.",
ctx.trace_id(),
timeout,
);
let (response_completion, response) = oneshot::channel();
let cancellation = self.cancellation.clone();
let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
Send {
fut: MapOkDispatchResponse::new(
MapErrConnectionReset::new(self.to_dispatch.send(DispatchRequest {
ctx,
request_id,
request,
response_completion,
})),
DispatchResponse {
response: tokio::time::timeout(timeout, response),
complete: false,
request_id,
cancellation,
ctx,
},
),
}
}
/// Sends a request to the dispatch task to forward to the server, returning a [`Future`] that
/// resolves to the response.
pub fn call(&mut self, context: context::Context, request: Req) -> Call<Req, Resp> {
Call {
fut: AndThenIdent::new(self.send(context, request)),
}
}
}
/// A server response that is completed by request dispatch when the corresponding response
/// arrives off the wire.
#[pin_project(PinnedDrop)]
#[derive(Debug)]
struct DispatchResponse<Resp> {
response: tokio::time::Timeout<oneshot::Receiver<Response<Resp>>>,
ctx: context::Context,
complete: bool,
cancellation: RequestCancellation,
request_id: u64,
}
impl<Resp> Future for DispatchResponse<Resp> {
type Output = io::Result<Resp>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<Resp>> {
let resp = ready!(self.response.poll_unpin(cx));
Poll::Ready(match resp {
Ok(resp) => {
self.complete = true;
match resp {
Ok(resp) => Ok(resp.message?),
Err(oneshot::Canceled) => {
// The oneshot is Canceled when the dispatch task ends. In that case,
// there's nothing listening on the other side, so there's no point in
// propagating cancellation.
Err(io::Error::from(io::ErrorKind::ConnectionReset))
}
}
}
Err(tokio::time::Elapsed { .. }) => Err(io::Error::new(
io::ErrorKind::TimedOut,
"Client dropped expired request.".to_string(),
)),
})
}
}
// Cancels the request when dropped, if not already complete.
#[pinned_drop]
impl<Resp> PinnedDrop for DispatchResponse<Resp> {
fn drop(mut self: Pin<&mut Self>) {
if !self.complete {
// The receiver needs to be closed to handle the edge case that the request has not
// yet been received by the dispatch task. It is possible for the cancel message to
// arrive before the request itself, in which case the request could get stuck in the
// dispatch map forever if the server never responds (e.g. if the server dies while
// responding). Even if the server does respond, it will have unnecessarily done work
// for a client no longer waiting for a response. To avoid this, the dispatch task
// checks if the receiver is closed before inserting the request in the map. By
// closing the receiver before sending the cancel message, it is guaranteed that if the
// dispatch task misses an early-arriving cancellation message, then it will see the
// receiver as closed.
self.response.get_mut().close();
let request_id = self.request_id;
self.cancellation.cancel(request_id);
}
}
}
/// Returns a channel and dispatcher that manages the lifecycle of requests initiated by the
/// channel.
pub fn new<Req, Resp, C>(
config: Config,
transport: C,
) -> NewClient<Channel<Req, Resp>, RequestDispatch<Req, Resp, C>>
where
C: Transport<ClientMessage<Req>, Response<Resp>>,
{
let (to_dispatch, pending_requests) = mpsc::channel(config.pending_request_buffer);
let (cancellation, canceled_requests) = cancellations();
let canceled_requests = canceled_requests.fuse();
NewClient {
client: Channel {
to_dispatch,
cancellation,
next_request_id: Arc::new(AtomicU64::new(0)),
},
dispatch: RequestDispatch {
config,
canceled_requests,
transport: transport.fuse(),
in_flight_requests: FnvHashMap::default(),
pending_requests: pending_requests.fuse(),
},
}
}
/// Handles the lifecycle of requests, writing requests to the wire, managing cancellations,
/// and dispatching responses to the appropriate channel.
#[pin_project]
#[derive(Debug)]
pub struct RequestDispatch<Req, Resp, C> {
/// Writes requests to the wire and reads responses off the wire.
#[pin]
transport: Fuse<C>,
/// Requests waiting to be written to the wire.
#[pin]
pending_requests: Fuse<mpsc::Receiver<DispatchRequest<Req, Resp>>>,
/// Requests that were dropped.
#[pin]
canceled_requests: Fuse<CanceledRequests>,
/// Requests already written to the wire that haven't yet received responses.
in_flight_requests: FnvHashMap<u64, InFlightData<Resp>>,
/// Configures limits to prevent unlimited resource usage.
config: Config,
}
impl<Req, Resp, C> RequestDispatch<Req, Resp, C>
where
C: Transport<ClientMessage<Req>, Response<Resp>>,
{
fn pump_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> PollIo<()> {
Poll::Ready(
match ready!(self.as_mut().project().transport.poll_next(cx)?) {
Some(response) => {
self.complete(response);
Some(Ok(()))
}
None => None,
},
)
}
fn pump_write(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> PollIo<()> {
enum ReceiverStatus {
NotReady,
Closed,
}
let pending_requests_status = match self.as_mut().poll_next_request(cx)? {
Poll::Ready(Some(dispatch_request)) => {
self.as_mut().write_request(dispatch_request)?;
return Poll::Ready(Some(Ok(())));
}
Poll::Ready(None) => ReceiverStatus::Closed,
Poll::Pending => ReceiverStatus::NotReady,
};
let canceled_requests_status = match self.as_mut().poll_next_cancellation(cx)? {
Poll::Ready(Some((context, request_id))) => {
self.as_mut().write_cancel(context, request_id)?;
return Poll::Ready(Some(Ok(())));
}
Poll::Ready(None) => ReceiverStatus::Closed,
Poll::Pending => ReceiverStatus::NotReady,
};
match (pending_requests_status, canceled_requests_status) {
(ReceiverStatus::Closed, ReceiverStatus::Closed) => {
ready!(self.as_mut().project().transport.poll_flush(cx)?);
Poll::Ready(None)
}
(ReceiverStatus::NotReady, _) | (_, ReceiverStatus::NotReady) => {
// No more messages to process, so flush any messages buffered in the transport.
ready!(self.as_mut().project().transport.poll_flush(cx)?);
// Even if we fully-flush, we return Pending, because we have no more requests
// or cancellations right now.
Poll::Pending
}
}
}
/// Yields the next pending request, if one is ready to be sent.
fn poll_next_request(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> PollIo<DispatchRequest<Req, Resp>> {
if self.as_mut().project().in_flight_requests.len() >= self.config.max_in_flight_requests {
info!(
"At in-flight request capacity ({}/{}).",
self.as_mut().project().in_flight_requests.len(),
self.config.max_in_flight_requests
);
// No need to schedule a wakeup, because timers and responses are responsible
// for clearing out in-flight requests.
return Poll::Pending;
}
while let Poll::Pending = self.as_mut().project().transport.poll_ready(cx)? {
// We can't yield a request-to-be-sent before the transport is capable of buffering it.
ready!(self.as_mut().project().transport.poll_flush(cx)?);
}
loop {
match ready!(self.as_mut().project().pending_requests.poll_next_unpin(cx)) {
Some(request) => {
if request.response_completion.is_canceled() {
trace!(
"[{}] Request canceled before being sent.",
request.ctx.trace_id()
);
continue;
}
return Poll::Ready(Some(Ok(request)));
}
None => return Poll::Ready(None),
}
}
}
/// Yields the next pending cancellation, and, if one is ready, cancels the associated request.
fn poll_next_cancellation(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> PollIo<(context::Context, u64)> {
while let Poll::Pending = self.as_mut().project().transport.poll_ready(cx)? {
ready!(self.as_mut().project().transport.poll_flush(cx)?);
}
loop {
let cancellation = self
.as_mut()
.project()
.canceled_requests
.poll_next_unpin(cx);
match ready!(cancellation) {
Some(request_id) => {
if let Some(in_flight_data) = self
.as_mut()
.project()
.in_flight_requests
.remove(&request_id)
{
self.as_mut().project().in_flight_requests.compact(0.1);
debug!("[{}] Removed request.", in_flight_data.ctx.trace_id());
return Poll::Ready(Some(Ok((in_flight_data.ctx, request_id))));
}
}
None => return Poll::Ready(None),
}
}
}
fn write_request(
mut self: Pin<&mut Self>,
dispatch_request: DispatchRequest<Req, Resp>,
) -> io::Result<()> {
let request_id = dispatch_request.request_id;
let request = ClientMessage::Request(Request {
id: request_id,
message: dispatch_request.request,
context: context::Context {
deadline: dispatch_request.ctx.deadline,
trace_context: dispatch_request.ctx.trace_context,
_non_exhaustive: (),
},
_non_exhaustive: (),
});
self.as_mut().project().transport.start_send(request)?;
self.as_mut().project().in_flight_requests.insert(
request_id,
InFlightData {
ctx: dispatch_request.ctx,
response_completion: dispatch_request.response_completion,
},
);
Ok(())
}
fn write_cancel(
mut self: Pin<&mut Self>,
context: context::Context,
request_id: u64,
) -> io::Result<()> {
let trace_id = *context.trace_id();
let cancel = ClientMessage::Cancel {
trace_context: context.trace_context,
request_id,
};
self.as_mut().project().transport.start_send(cancel)?;
trace!("[{}] Cancel message sent.", trace_id);
Ok(())
}
/// Sends a server response to the client task that initiated the associated request.
fn complete(mut self: Pin<&mut Self>, response: Response<Resp>) -> bool {
if let Some(in_flight_data) = self
.as_mut()
.project()
.in_flight_requests
.remove(&response.request_id)
{
self.as_mut().project().in_flight_requests.compact(0.1);
trace!("[{}] Received response.", in_flight_data.ctx.trace_id());
let _ = in_flight_data.response_completion.send(response);
return true;
}
debug!(
"No in-flight request found for request_id = {}.",
response.request_id
);
// If the response completion was absent, then the request was already canceled.
false
}
}
impl<Req, Resp, C> Future for RequestDispatch<Req, Resp, C>
where
C: Transport<ClientMessage<Req>, Response<Resp>>,
{
type Output = io::Result<()>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
loop {
match (self.as_mut().pump_read(cx)?, self.as_mut().pump_write(cx)?) {
(read, Poll::Ready(None)) => {
if self.as_mut().project().in_flight_requests.is_empty() {
info!("Shutdown: write half closed, and no requests in flight.");
return Poll::Ready(Ok(()));
}
info!(
"Shutdown: write half closed, and {} requests in flight.",
self.as_mut().project().in_flight_requests.len()
);
match read {
Poll::Ready(Some(())) => continue,
_ => return Poll::Pending,
}
}
(Poll::Ready(Some(())), _) | (_, Poll::Ready(Some(()))) => {}
_ => return Poll::Pending,
}
}
}
}
/// A server-bound request sent from a [`Channel`] to request dispatch, which will then manage
/// the lifecycle of the request.
#[derive(Debug)]
struct DispatchRequest<Req, Resp> {
ctx: context::Context,
request_id: u64,
request: Req,
response_completion: oneshot::Sender<Response<Resp>>,
}
#[derive(Debug)]
struct InFlightData<Resp> {
ctx: context::Context,
response_completion: oneshot::Sender<Response<Resp>>,
}
/// Sends request cancellation signals.
#[derive(Debug, Clone)]
struct RequestCancellation(mpsc::UnboundedSender<u64>);
/// A stream of IDs of requests that have been canceled.
#[derive(Debug)]
struct CanceledRequests(mpsc::UnboundedReceiver<u64>);
/// Returns a channel to send request cancellation messages.
fn cancellations() -> (RequestCancellation, CanceledRequests) {
// Unbounded because messages are sent in the drop fn. This is fine, because it's still
// bounded by the number of in-flight requests. Additionally, each request has a clone
// of the sender, so the bounded channel would have the same behavior,
// since it guarantees a slot.
let (tx, rx) = mpsc::unbounded();
(RequestCancellation(tx), CanceledRequests(rx))
}
impl RequestCancellation {
/// Cancels the request with ID `request_id`.
fn cancel(&mut self, request_id: u64) {
let _ = self.0.unbounded_send(request_id);
}
}
impl Stream for CanceledRequests {
type Item = u64;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<u64>> {
self.0.poll_next_unpin(cx)
}
}
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
struct MapErrConnectionReset<Fut> {
#[pin]
future: Fut,
finished: Option<()>,
}
impl<Fut> MapErrConnectionReset<Fut> {
fn new(future: Fut) -> MapErrConnectionReset<Fut> {
MapErrConnectionReset {
future,
finished: Some(()),
}
}
}
impl<Fut> Future for MapErrConnectionReset<Fut>
where
Fut: TryFuture,
{
type Output = io::Result<Fut::Ok>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.as_mut().project().future.try_poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(result) => {
self.project().finished.take().expect(
"MapErrConnectionReset must not be polled after it returned `Poll::Ready`",
);
Poll::Ready(result.map_err(|_| io::Error::from(io::ErrorKind::ConnectionReset)))
}
}
}
}
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
struct MapOkDispatchResponse<Fut, Resp> {
#[pin]
future: Fut,
response: Option<DispatchResponse<Resp>>,
}
impl<Fut, Resp> MapOkDispatchResponse<Fut, Resp> {
fn new(future: Fut, response: DispatchResponse<Resp>) -> MapOkDispatchResponse<Fut, Resp> {
MapOkDispatchResponse {
future,
response: Some(response),
}
}
}
impl<Fut, Resp> Future for MapOkDispatchResponse<Fut, Resp>
where
Fut: TryFuture,
{
type Output = Result<DispatchResponse<Resp>, Fut::Error>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.as_mut().project().future.try_poll(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(result) => {
let response = self
.as_mut()
.project()
.response
.take()
.expect("MapOk must not be polled after it returned `Poll::Ready`");
Poll::Ready(result.map(|_| response))
}
}
}
}
#[pin_project]
#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
struct AndThenIdent<Fut1, Fut2> {
#[pin]
try_chain: TryChain<Fut1, Fut2>,
}
impl<Fut1, Fut2> AndThenIdent<Fut1, Fut2>
where
Fut1: TryFuture<Ok = Fut2>,
Fut2: TryFuture,
{
/// Creates a new `Then`.
fn new(future: Fut1) -> AndThenIdent<Fut1, Fut2> {
AndThenIdent {
try_chain: TryChain::new(future),
}
}
}
impl<Fut1, Fut2> Future for AndThenIdent<Fut1, Fut2>
where
Fut1: TryFuture<Ok = Fut2>,
Fut2: TryFuture<Error = Fut1::Error>,
{
type Output = Result<Fut2::Ok, Fut2::Error>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
self.project().try_chain.poll(cx, |result| match result {
Ok(ok) => TryChainAction::Future(ok),
Err(err) => TryChainAction::Output(Err(err)),
})
}
}
#[must_use = "futures do nothing unless polled"]
#[derive(Debug)]
enum TryChain<Fut1, Fut2> {
First(Fut1),
Second(Fut2),
Empty,
}
enum TryChainAction<Fut2>
where
Fut2: TryFuture,
{
Future(Fut2),
Output(Result<Fut2::Ok, Fut2::Error>),
}
impl<Fut1, Fut2> TryChain<Fut1, Fut2>
where
Fut1: TryFuture<Ok = Fut2>,
Fut2: TryFuture,
{
fn new(fut1: Fut1) -> TryChain<Fut1, Fut2> {
TryChain::First(fut1)
}
fn poll<F>(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
f: F,
) -> Poll<Result<Fut2::Ok, Fut2::Error>>
where
F: FnOnce(Result<Fut1::Ok, Fut1::Error>) -> TryChainAction<Fut2>,
{
let mut f = Some(f);
// Safe to call `get_unchecked_mut` because we won't move the futures.
let this = unsafe { Pin::get_unchecked_mut(self) };
loop {
let output = match this {
TryChain::First(fut1) => {
// Poll the first future
match unsafe { Pin::new_unchecked(fut1) }.try_poll(cx) {
Poll::Pending => return Poll::Pending,
Poll::Ready(output) => output,
}
}
TryChain::Second(fut2) => {
// Poll the second future
return unsafe { Pin::new_unchecked(fut2) }.try_poll(cx);
}
TryChain::Empty => {
panic!("future must not be polled after it returned `Poll::Ready`");
}
};
*this = TryChain::Empty; // Drop fut1
let f = f.take().unwrap();
match f(output) {
TryChainAction::Future(fut2) => *this = TryChain::Second(fut2),
TryChainAction::Output(output) => return Poll::Ready(output),
}
}
}
}
#[cfg(test)]
mod tests {
use super::{
cancellations, CanceledRequests, Channel, DispatchResponse, RequestCancellation,
RequestDispatch,
};
use crate::{
client::Config,
context,
transport::{self, channel::UnboundedChannel},
ClientMessage, Response,
};
use fnv::FnvHashMap;
use futures::{
channel::{mpsc, oneshot},
prelude::*,
task::*,
};
use std::time::Duration;
use std::{pin::Pin, sync::atomic::AtomicU64, sync::Arc};
#[tokio::test(threaded_scheduler)]
async fn dispatch_response_cancels_on_timeout() {
let (_response_completion, response) = oneshot::channel();
let (cancellation, mut canceled_requests) = cancellations();
let resp = DispatchResponse::<u64> {
// Timeout in the past should cause resp to error out when polled.
response: tokio::time::timeout(Duration::from_secs(0), response),
complete: false,
request_id: 3,
cancellation,
ctx: context::current(),
};
let _ = futures::poll!(resp);
// resp's drop() is run, which should send a cancel message.
assert!(canceled_requests.0.try_next().unwrap() == Some(3));
}
#[tokio::test(threaded_scheduler)]
async fn stage_request() {
let (mut dispatch, mut channel, _server_channel) = set_up();
let dispatch = Pin::new(&mut dispatch);
let cx = &mut Context::from_waker(&noop_waker_ref());
let _resp = send_request(&mut channel, "hi").await;
let req = dispatch.poll_next_request(cx).ready();
assert!(req.is_some());
let req = req.unwrap();
assert_eq!(req.request_id, 0);
assert_eq!(req.request, "hi".to_string());
}
// Regression test for https://github.com/google/tarpc/issues/220
#[tokio::test(threaded_scheduler)]
async fn stage_request_channel_dropped_doesnt_panic() {
let (mut dispatch, mut channel, mut server_channel) = set_up();
let mut dispatch = Pin::new(&mut dispatch);
let cx = &mut Context::from_waker(&noop_waker_ref());
let _ = send_request(&mut channel, "hi").await;
drop(channel);
assert!(dispatch.as_mut().poll(cx).is_ready());
send_response(
&mut server_channel,
Response {
request_id: 0,
message: Ok("hello".into()),
_non_exhaustive: (),
},
)
.await;
dispatch.await.unwrap();
}
#[tokio::test(threaded_scheduler)]
async fn stage_request_response_future_dropped_is_canceled_before_sending() {
let (mut dispatch, mut channel, _server_channel) = set_up();
let dispatch = Pin::new(&mut dispatch);
let cx = &mut Context::from_waker(&noop_waker_ref());
let _ = send_request(&mut channel, "hi").await;
// Drop the channel so polling returns none if no requests are currently ready.
drop(channel);
// Test that a request future dropped before it's processed by dispatch will cause the request
// to not be added to the in-flight request map.
assert!(dispatch.poll_next_request(cx).ready().is_none());
}
#[tokio::test(threaded_scheduler)]
async fn stage_request_response_future_dropped_is_canceled_after_sending() {
let (mut dispatch, mut channel, _server_channel) = set_up();
let cx = &mut Context::from_waker(&noop_waker_ref());
let mut dispatch = Pin::new(&mut dispatch);
let req = send_request(&mut channel, "hi").await;
assert!(dispatch.as_mut().pump_write(cx).ready().is_some());
assert!(!dispatch.as_mut().project().in_flight_requests.is_empty());
// Test that a request future dropped after it's processed by dispatch will cause the request
// to be removed from the in-flight request map.
drop(req);
if let Poll::Ready(Some(_)) = dispatch.as_mut().poll_next_cancellation(cx).unwrap() {
// ok
} else {
panic!("Expected request to be cancelled")
};
assert!(dispatch.project().in_flight_requests.is_empty());
}
#[tokio::test(threaded_scheduler)]
async fn stage_request_response_closed_skipped() {
let (mut dispatch, mut channel, _server_channel) = set_up();
let dispatch = Pin::new(&mut dispatch);
let cx = &mut Context::from_waker(&noop_waker_ref());
// Test that a request future that's closed its receiver but not yet canceled its request --
// i.e. still in `drop fn` -- will cause the request to not be added to the in-flight request
// map.
let mut resp = send_request(&mut channel, "hi").await;
resp.response.get_mut().close();
assert!(dispatch.poll_next_request(cx).is_pending());
}
fn set_up() -> (
RequestDispatch<String, String, UnboundedChannel<Response<String>, ClientMessage<String>>>,
Channel<String, String>,
UnboundedChannel<ClientMessage<String>, Response<String>>,
) {
let _ = env_logger::try_init();
let (to_dispatch, pending_requests) = mpsc::channel(1);
let (cancel_tx, canceled_requests) = mpsc::unbounded();
let (client_channel, server_channel) = transport::channel::unbounded();
let dispatch = RequestDispatch::<String, String, _> {
transport: client_channel.fuse(),
pending_requests: pending_requests.fuse(),
canceled_requests: CanceledRequests(canceled_requests).fuse(),
in_flight_requests: FnvHashMap::default(),
config: Config::default(),
};
let cancellation = RequestCancellation(cancel_tx);
let channel = Channel {
to_dispatch,
cancellation,
next_request_id: Arc::new(AtomicU64::new(0)),
};
(dispatch, channel, server_channel)
}
async fn send_request(
channel: &mut Channel<String, String>,
request: &str,
) -> DispatchResponse<String> {
channel
.send(context::current(), request.to_string())
.await
.unwrap()
}
async fn send_response(
channel: &mut UnboundedChannel<ClientMessage<String>, Response<String>>,
response: Response<String>,
) {
channel.send(response).await.unwrap();
}
trait PollTest {
type T;
fn unwrap(self) -> Poll<Self::T>;
fn ready(self) -> Self::T;
}
impl<T, E> PollTest for Poll<Option<Result<T, E>>>
where
E: ::std::fmt::Display,
{
type T = Option<T>;
fn unwrap(self) -> Poll<Option<T>> {
match self {
Poll::Ready(Some(Ok(t))) => Poll::Ready(Some(t)),
Poll::Ready(None) => Poll::Ready(None),
Poll::Ready(Some(Err(e))) => panic!(e.to_string()),
Poll::Pending => Poll::Pending,
}
}
fn ready(self) -> Option<T> {
match self {
Poll::Ready(Some(Ok(t))) => Some(t),
Poll::Ready(None) => None,
Poll::Ready(Some(Err(e))) => panic!(e.to_string()),
Poll::Pending => panic!("Pending"),
}
}
}
}

View File

@@ -1,155 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a client that connects to a server and sends multiplexed requests.
use crate::context;
use futures::prelude::*;
use std::io;
/// Provides a [`Client`] backed by a transport.
pub mod channel;
pub use channel::{new, Channel};
/// Sends multiplexed requests to, and receives responses from, a server.
pub trait Client<'a, Req> {
/// The response type.
type Response;
/// The future response.
type Future: Future<Output = io::Result<Self::Response>> + 'a;
/// Initiates a request, sending it to the dispatch task.
///
/// Returns a [`Future`] that resolves to this client and the future response
/// once the request is successfully enqueued.
///
/// [`Future`]: futures::Future
fn call(&'a mut self, ctx: context::Context, request: Req) -> Self::Future;
/// Returns a Client that applies a post-processing function to the returned response.
fn map_response<F, R>(self, f: F) -> MapResponse<Self, F>
where
F: FnMut(Self::Response) -> R,
Self: Sized,
{
MapResponse { inner: self, f }
}
/// Returns a Client that applies a pre-processing function to the request.
fn with_request<F, Req2>(self, f: F) -> WithRequest<Self, F>
where
F: FnMut(Req2) -> Req,
Self: Sized,
{
WithRequest { inner: self, f }
}
}
/// A Client that applies a function to the returned response.
#[derive(Clone, Debug)]
pub struct MapResponse<C, F> {
inner: C,
f: F,
}
impl<'a, C, F, Req, Resp, Resp2> Client<'a, Req> for MapResponse<C, F>
where
C: Client<'a, Req, Response = Resp>,
F: FnMut(Resp) -> Resp2 + 'a,
{
type Response = Resp2;
type Future = futures::future::MapOk<<C as Client<'a, Req>>::Future, &'a mut F>;
fn call(&'a mut self, ctx: context::Context, request: Req) -> Self::Future {
self.inner.call(ctx, request).map_ok(&mut self.f)
}
}
/// A Client that applies a pre-processing function to the request.
#[derive(Clone, Debug)]
pub struct WithRequest<C, F> {
inner: C,
f: F,
}
impl<'a, C, F, Req, Req2, Resp> Client<'a, Req2> for WithRequest<C, F>
where
C: Client<'a, Req, Response = Resp>,
F: FnMut(Req2) -> Req,
{
type Response = Resp;
type Future = <C as Client<'a, Req>>::Future;
fn call(&'a mut self, ctx: context::Context, request: Req2) -> Self::Future {
self.inner.call(ctx, (self.f)(request))
}
}
impl<'a, Req, Resp> Client<'a, Req> for Channel<Req, Resp>
where
Req: 'a,
Resp: 'a,
{
type Response = Resp;
type Future = channel::Call<'a, Req, Resp>;
fn call(&'a mut self, ctx: context::Context, request: Req) -> channel::Call<'a, Req, Resp> {
self.call(ctx, request)
}
}
/// Settings that control the behavior of the client.
#[derive(Clone, Debug)]
pub struct Config {
/// The number of requests that can be in flight at once.
/// `max_in_flight_requests` controls the size of the map used by the client
/// for storing pending requests.
pub max_in_flight_requests: usize,
/// The number of requests that can be buffered client-side before being sent.
/// `pending_requests_buffer` controls the size of the channel clients use
/// to communicate with the request dispatch task.
pub pending_request_buffer: usize,
#[doc(hidden)]
_non_exhaustive: (),
}
impl Default for Config {
fn default() -> Self {
Config {
max_in_flight_requests: 1_000,
pending_request_buffer: 100,
_non_exhaustive: (),
}
}
}
/// A channel and dispatch pair. The dispatch drives the sending and receiving of requests
/// and must be polled continuously or spawned.
#[derive(Debug)]
pub struct NewClient<C, D> {
/// The new client.
pub client: C,
/// The client's dispatch.
pub dispatch: D,
}
impl<C, D> NewClient<C, D>
where
D: Future<Output = io::Result<()>> + Send + 'static,
{
/// Helper method to spawn the dispatch on the default executor.
#[cfg(feature = "tokio1")]
pub fn spawn(self) -> io::Result<C> {
use log::error;
let dispatch = self
.dispatch
.unwrap_or_else(move |e| error!("Connection broken: {}", e));
tokio::spawn(dispatch);
Ok(self.client)
}
}

View File

@@ -1,63 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a request context that carries a deadline and trace context. This context is sent from
//! client to server and is used by the server to enforce response deadlines.
use crate::trace::{self, TraceId};
use std::time::{Duration, SystemTime};
/// A request context that carries request-scoped information like deadlines and trace information.
/// It is sent from client to server and is used by the server to enforce response deadlines.
///
/// The context should not be stored directly in a server implementation, because the context will
/// be different for each request in scope.
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Context {
/// When the client expects the request to be complete by. The server should cancel the request
/// if it is not complete by this time.
#[cfg_attr(
feature = "serde1",
serde(serialize_with = "crate::util::serde::serialize_epoch_secs")
)]
#[cfg_attr(
feature = "serde1",
serde(deserialize_with = "crate::util::serde::deserialize_epoch_secs")
)]
#[cfg_attr(feature = "serde1", serde(default = "ten_seconds_from_now"))]
pub deadline: SystemTime,
/// Uniquely identifies requests originating from the same source.
/// When a service handles a request by making requests itself, those requests should
/// include the same `trace_id` as that included on the original request. This way,
/// users can trace related actions across a distributed system.
pub trace_context: trace::Context,
#[doc(hidden)]
#[cfg_attr(feature = "serde1", serde(skip_serializing, default))]
pub(crate) _non_exhaustive: (),
}
#[cfg(feature = "serde1")]
fn ten_seconds_from_now() -> SystemTime {
SystemTime::now() + Duration::from_secs(10)
}
/// Returns the context for the current request, or a default Context if no request is active.
// TODO: populate Context with request-scoped data, with default fallbacks.
pub fn current() -> Context {
Context {
deadline: SystemTime::now() + Duration::from_secs(10),
trace_context: trace::Context::new_root(),
_non_exhaustive: (),
}
}
impl Context {
/// Returns the ID of the request-scoped trace.
pub fn trace_id(&self) -> &TraceId {
&self.trace_context.trace_id
}
}

View File

@@ -1,127 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
#![deny(missing_docs, missing_debug_implementations)]
//! An RPC framework providing client and server.
//!
//! Features:
//! * RPC deadlines, both client- and server-side.
//! * Cascading cancellation (works with multiple hops).
//! * Configurable limits
//! * In-flight requests, both client and server-side.
//! * Server-side limit is per-connection.
//! * When the server reaches the in-flight request maximum, it returns a throttled error
//! to the client.
//! * When the client reaches the in-flight request max, messages are buffered up to a
//! configurable maximum, beyond which the requests are back-pressured.
//! * Server connections.
//! * Total and per-IP limits.
//! * When an incoming connection is accepted, if already at maximum, the connection is
//! dropped.
//! * Transport agnostic.
pub mod client;
pub mod context;
pub mod server;
pub mod transport;
pub(crate) mod util;
pub use crate::{client::Client, server::Server, trace, transport::sealed::Transport};
use futures::task::*;
use std::{io, time::SystemTime};
/// A message from a client to a server.
#[derive(Debug)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub enum ClientMessage<T> {
/// A request initiated by a user. The server responds to a request by invoking a
/// service-provided request handler. The handler completes with a [`response`](Response), which
/// the server sends back to the client.
Request(Request<T>),
/// A command to cancel an in-flight request, automatically sent by the client when a response
/// future is dropped.
///
/// When received, the server will immediately cancel the main task (top-level future) of the
/// request handler for the associated request. Any tasks spawned by the request handler will
/// not be canceled, because the framework layer does not
/// know about them.
Cancel {
/// The trace context associates the message with a specific chain of causally-related actions,
/// possibly orchestrated across many distributed systems.
#[cfg_attr(feature = "serde", serde(default))]
trace_context: trace::Context,
/// The ID of the request to cancel.
request_id: u64,
},
#[doc(hidden)]
_NonExhaustive,
}
/// A request from a client to a server.
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Request<T> {
/// Trace context, deadline, and other cross-cutting concerns.
pub context: context::Context,
/// Uniquely identifies the request across all requests sent over a single channel.
pub id: u64,
/// The request body.
pub message: T,
#[doc(hidden)]
#[cfg_attr(feature = "serde1", serde(skip_serializing, default))]
_non_exhaustive: (),
}
/// A response from a server to a client.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Response<T> {
/// The ID of the request being responded to.
pub request_id: u64,
/// The response body, or an error if the request failed.
pub message: Result<T, ServerError>,
#[doc(hidden)]
#[cfg_attr(feature = "serde1", serde(skip_serializing, default))]
_non_exhaustive: (),
}
/// An error response from a server to a client.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct ServerError {
#[cfg_attr(
feature = "serde1",
serde(serialize_with = "util::serde::serialize_io_error_kind_as_u32")
)]
#[cfg_attr(
feature = "serde1",
serde(deserialize_with = "util::serde::deserialize_io_error_kind_from_u32")
)]
/// The type of error that occurred to fail the request.
pub kind: io::ErrorKind,
/// A message describing more detail about the error that occurred.
pub detail: Option<String>,
#[doc(hidden)]
#[cfg_attr(feature = "serde1", serde(skip_serializing, default))]
_non_exhaustive: (),
}
impl From<ServerError> for io::Error {
fn from(e: ServerError) -> io::Error {
io::Error::new(e.kind, e.detail.unwrap_or_default())
}
}
impl<T> Request<T> {
/// Returns the deadline for this request.
pub fn deadline(&self) -> &SystemTime {
&self.context.deadline
}
}
pub(crate) type PollIo<T> = Poll<Option<io::Result<T>>>;

View File

@@ -1,702 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a server that concurrently handles many connections sending multiplexed requests.
use crate::{
context, trace, util::Compact, util::TimeUntil, ClientMessage, PollIo, Request, Response,
ServerError, Transport,
};
use fnv::FnvHashMap;
use futures::{
channel::mpsc,
future::{AbortHandle, AbortRegistration, Abortable},
prelude::*,
ready,
stream::Fuse,
task::*,
};
use humantime::format_rfc3339;
use log::{debug, trace};
use pin_project::pin_project;
use std::{fmt, hash::Hash, io, marker::PhantomData, pin::Pin, time::SystemTime};
use tokio::time::Timeout;
mod filter;
#[cfg(test)]
mod testing;
mod throttle;
pub use self::{
filter::ChannelFilter,
throttle::{Throttler, ThrottlerStream},
};
/// Manages clients, serving multiplexed requests over each connection.
#[derive(Debug)]
pub struct Server<Req, Resp> {
config: Config,
ghost: PhantomData<(Req, Resp)>,
}
impl<Req, Resp> Default for Server<Req, Resp> {
fn default() -> Self {
new(Config::default())
}
}
/// Settings that control the behavior of the server.
#[derive(Clone, Debug)]
pub struct Config {
/// The number of responses per client that can be buffered server-side before being sent.
/// `pending_response_buffer` controls the buffer size of the channel that a server's
/// response tasks use to send responses to the client handler task.
pub pending_response_buffer: usize,
}
impl Default for Config {
fn default() -> Self {
Config {
pending_response_buffer: 100,
}
}
}
impl Config {
/// Returns a channel backed by `transport` and configured with `self`.
pub fn channel<Req, Resp, T>(self, transport: T) -> BaseChannel<Req, Resp, T>
where
T: Transport<Response<Resp>, ClientMessage<Req>>,
{
BaseChannel::new(self, transport)
}
}
/// Returns a new server with configuration specified `config`.
pub fn new<Req, Resp>(config: Config) -> Server<Req, Resp> {
Server {
config,
ghost: PhantomData,
}
}
impl<Req, Resp> Server<Req, Resp> {
/// Returns the config for this server.
pub fn config(&self) -> &Config {
&self.config
}
/// Returns a stream of server channels.
pub fn incoming<S, T>(self, listener: S) -> impl Stream<Item = BaseChannel<Req, Resp, T>>
where
S: Stream<Item = T>,
T: Transport<Response<Resp>, ClientMessage<Req>>,
{
listener.map(move |t| BaseChannel::new(self.config.clone(), t))
}
}
/// Basically a Fn(Req) -> impl Future<Output = Resp>;
pub trait Serve<Req>: Sized + Clone {
/// Type of response.
type Resp;
/// Type of response future.
type Fut: Future<Output = Self::Resp>;
/// Responds to a single request.
fn serve(self, ctx: context::Context, req: Req) -> Self::Fut;
}
impl<Req, Resp, Fut, F> Serve<Req> for F
where
F: FnOnce(context::Context, Req) -> Fut + Clone,
Fut: Future<Output = Resp>,
{
type Resp = Resp;
type Fut = Fut;
fn serve(self, ctx: context::Context, req: Req) -> Self::Fut {
self(ctx, req)
}
}
/// A utility trait enabling a stream to fluently chain a request handler.
pub trait Handler<C>
where
Self: Sized + Stream<Item = C>,
C: Channel,
{
/// Enforces channel per-key limits.
fn max_channels_per_key<K, KF>(self, n: u32, keymaker: KF) -> filter::ChannelFilter<Self, K, KF>
where
K: fmt::Display + Eq + Hash + Clone + Unpin,
KF: Fn(&C) -> K,
{
ChannelFilter::new(self, n, keymaker)
}
/// Caps the number of concurrent requests per channel.
fn max_concurrent_requests_per_channel(self, n: usize) -> ThrottlerStream<Self> {
ThrottlerStream::new(self, n)
}
/// Responds to all requests with `server`.
#[cfg(feature = "tokio1")]
fn respond_with<S>(self, server: S) -> Running<Self, S>
where
S: Serve<C::Req, Resp = C::Resp>,
{
Running {
incoming: self,
server,
}
}
}
impl<S, C> Handler<C> for S
where
S: Sized + Stream<Item = C>,
C: Channel,
{
}
/// BaseChannel lifts a Transport to a Channel by tracking in-flight requests.
#[pin_project]
#[derive(Debug)]
pub struct BaseChannel<Req, Resp, T> {
config: Config,
/// Writes responses to the wire and reads requests off the wire.
#[pin]
transport: Fuse<T>,
/// Number of requests currently being responded to.
in_flight_requests: FnvHashMap<u64, AbortHandle>,
/// Types the request and response.
ghost: PhantomData<(Req, Resp)>,
}
impl<Req, Resp, T> BaseChannel<Req, Resp, T>
where
T: Transport<Response<Resp>, ClientMessage<Req>>,
{
/// Creates a new channel backed by `transport` and configured with `config`.
pub fn new(config: Config, transport: T) -> Self {
BaseChannel {
config,
transport: transport.fuse(),
in_flight_requests: FnvHashMap::default(),
ghost: PhantomData,
}
}
/// Creates a new channel backed by `transport` and configured with the defaults.
pub fn with_defaults(transport: T) -> Self {
Self::new(Config::default(), transport)
}
/// Returns the inner transport.
pub fn get_ref(&self) -> &T {
self.transport.get_ref()
}
fn cancel_request(mut self: Pin<&mut Self>, trace_context: &trace::Context, request_id: u64) {
// It's possible the request was already completed, so it's fine
// if this is None.
if let Some(cancel_handle) = self
.as_mut()
.project()
.in_flight_requests
.remove(&request_id)
{
self.as_mut().project().in_flight_requests.compact(0.1);
cancel_handle.abort();
let remaining = self.as_mut().project().in_flight_requests.len();
trace!(
"[{}] Request canceled. In-flight requests = {}",
trace_context.trace_id,
remaining,
);
} else {
trace!(
"[{}] Received cancellation, but response handler \
is already complete.",
trace_context.trace_id,
);
}
}
}
/// The server end of an open connection with a client, streaming in requests from, and sinking
/// responses to, the client.
///
/// Channels are free to somewhat rely on the assumption that all in-flight requests are eventually
/// either [cancelled](BaseChannel::cancel_request) or [responded to](Sink::start_send). Safety cannot
/// rely on this assumption, but it is best for `Channel` users to always account for all outstanding
/// requests.
pub trait Channel
where
Self: Transport<Response<<Self as Channel>::Resp>, Request<<Self as Channel>::Req>>,
{
/// Type of request item.
type Req;
/// Type of response sink item.
type Resp;
/// Configuration of the channel.
fn config(&self) -> &Config;
/// Returns the number of in-flight requests over this channel.
fn in_flight_requests(self: Pin<&mut Self>) -> usize;
/// Caps the number of concurrent requests.
fn max_concurrent_requests(self, n: usize) -> Throttler<Self>
where
Self: Sized,
{
Throttler::new(self, n)
}
/// Tells the Channel that request with ID `request_id` is being handled.
/// The request will be tracked until a response with the same ID is sent
/// to the Channel.
fn start_request(self: Pin<&mut Self>, request_id: u64) -> AbortRegistration;
/// Respond to requests coming over the channel with `f`. Returns a future that drives the
/// responses and resolves when the connection is closed.
fn respond_with<S>(self, server: S) -> ClientHandler<Self, S>
where
S: Serve<Self::Req, Resp = Self::Resp>,
Self: Sized,
{
let (responses_tx, responses) = mpsc::channel(self.config().pending_response_buffer);
let responses = responses.fuse();
ClientHandler {
channel: self,
server,
pending_responses: responses,
responses_tx,
}
}
}
impl<Req, Resp, T> Stream for BaseChannel<Req, Resp, T>
where
T: Transport<Response<Resp>, ClientMessage<Req>>,
{
type Item = io::Result<Request<Req>>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
loop {
match ready!(self.as_mut().project().transport.poll_next(cx)?) {
Some(message) => match message {
ClientMessage::Request(request) => {
return Poll::Ready(Some(Ok(request)));
}
ClientMessage::Cancel {
trace_context,
request_id,
} => {
self.as_mut().cancel_request(&trace_context, request_id);
}
ClientMessage::_NonExhaustive => unreachable!(),
},
None => return Poll::Ready(None),
}
}
}
}
impl<Req, Resp, T> Sink<Response<Resp>> for BaseChannel<Req, Resp, T>
where
T: Transport<Response<Resp>, ClientMessage<Req>>,
{
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().transport.poll_ready(cx)
}
fn start_send(mut self: Pin<&mut Self>, response: Response<Resp>) -> Result<(), Self::Error> {
if self
.as_mut()
.project()
.in_flight_requests
.remove(&response.request_id)
.is_some()
{
self.as_mut().project().in_flight_requests.compact(0.1);
}
self.project().transport.start_send(response)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().transport.poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().transport.poll_close(cx)
}
}
impl<Req, Resp, T> AsRef<T> for BaseChannel<Req, Resp, T> {
fn as_ref(&self) -> &T {
self.transport.get_ref()
}
}
impl<Req, Resp, T> Channel for BaseChannel<Req, Resp, T>
where
T: Transport<Response<Resp>, ClientMessage<Req>>,
{
type Req = Req;
type Resp = Resp;
fn config(&self) -> &Config {
&self.config
}
fn in_flight_requests(mut self: Pin<&mut Self>) -> usize {
self.as_mut().project().in_flight_requests.len()
}
fn start_request(self: Pin<&mut Self>, request_id: u64) -> AbortRegistration {
let (abort_handle, abort_registration) = AbortHandle::new_pair();
assert!(self
.project()
.in_flight_requests
.insert(request_id, abort_handle)
.is_none());
abort_registration
}
}
/// A running handler serving all requests coming over a channel.
#[pin_project]
#[derive(Debug)]
pub struct ClientHandler<C, S>
where
C: Channel,
{
#[pin]
channel: C,
/// Responses waiting to be written to the wire.
#[pin]
pending_responses: Fuse<mpsc::Receiver<(context::Context, Response<C::Resp>)>>,
/// Handed out to request handlers to fan in responses.
#[pin]
responses_tx: mpsc::Sender<(context::Context, Response<C::Resp>)>,
/// Server
server: S,
}
impl<C, S> ClientHandler<C, S>
where
C: Channel,
S: Serve<C::Req, Resp = C::Resp>,
{
fn pump_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> PollIo<RequestHandler<S::Fut, C::Resp>> {
match ready!(self.as_mut().project().channel.poll_next(cx)?) {
Some(request) => Poll::Ready(Some(Ok(self.handle_request(request)))),
None => Poll::Ready(None),
}
}
fn pump_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
read_half_closed: bool,
) -> PollIo<()> {
match self.as_mut().poll_next_response(cx)? {
Poll::Ready(Some((ctx, response))) => {
trace!(
"[{}] Staging response. In-flight requests = {}.",
ctx.trace_id(),
self.as_mut().project().channel.in_flight_requests(),
);
self.as_mut().project().channel.start_send(response)?;
Poll::Ready(Some(Ok(())))
}
Poll::Ready(None) => {
// Shutdown can't be done before we finish pumping out remaining responses.
ready!(self.as_mut().project().channel.poll_flush(cx)?);
Poll::Ready(None)
}
Poll::Pending => {
// No more requests to process, so flush any requests buffered in the transport.
ready!(self.as_mut().project().channel.poll_flush(cx)?);
// Being here means there are no staged requests and all written responses are
// fully flushed. So, if the read half is closed and there are no in-flight
// requests, then we can close the write half.
if read_half_closed && self.as_mut().project().channel.in_flight_requests() == 0 {
Poll::Ready(None)
} else {
Poll::Pending
}
}
}
}
fn poll_next_response(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> PollIo<(context::Context, Response<C::Resp>)> {
// Ensure there's room to write a response.
while let Poll::Pending = self.as_mut().project().channel.poll_ready(cx)? {
ready!(self.as_mut().project().channel.poll_flush(cx)?);
}
match ready!(self.as_mut().project().pending_responses.poll_next(cx)) {
Some((ctx, response)) => Poll::Ready(Some(Ok((ctx, response)))),
None => {
// This branch likely won't happen, since the ClientHandler is holding a Sender.
Poll::Ready(None)
}
}
}
fn handle_request(
mut self: Pin<&mut Self>,
request: Request<C::Req>,
) -> RequestHandler<S::Fut, C::Resp> {
let request_id = request.id;
let deadline = request.context.deadline;
let timeout = deadline.time_until();
trace!(
"[{}] Received request with deadline {} (timeout {:?}).",
request.context.trace_id(),
format_rfc3339(deadline),
timeout,
);
let ctx = request.context;
let request = request.message;
let response = self.as_mut().project().server.clone().serve(ctx, request);
let response = Resp {
state: RespState::PollResp,
request_id,
ctx,
deadline,
f: tokio::time::timeout(timeout, response),
response: None,
response_tx: self.as_mut().project().responses_tx.clone(),
};
let abort_registration = self.as_mut().project().channel.start_request(request_id);
RequestHandler {
resp: Abortable::new(response, abort_registration),
}
}
}
/// A future fulfilling a single client request.
#[pin_project]
#[derive(Debug)]
pub struct RequestHandler<F, R> {
#[pin]
resp: Abortable<Resp<F, R>>,
}
impl<F, R> Future for RequestHandler<F, R>
where
F: Future<Output = R>,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let _ = ready!(self.project().resp.poll(cx));
Poll::Ready(())
}
}
#[pin_project]
#[derive(Debug)]
struct Resp<F, R> {
state: RespState,
request_id: u64,
ctx: context::Context,
deadline: SystemTime,
#[pin]
f: Timeout<F>,
response: Option<Response<R>>,
#[pin]
response_tx: mpsc::Sender<(context::Context, Response<R>)>,
}
#[derive(Debug)]
#[allow(clippy::enum_variant_names)]
enum RespState {
PollResp,
PollReady,
PollFlush,
}
impl<F, R> Future for Resp<F, R>
where
F: Future<Output = R>,
{
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
loop {
match self.as_mut().project().state {
RespState::PollResp => {
let result = ready!(self.as_mut().project().f.poll(cx));
*self.as_mut().project().response = Some(Response {
request_id: self.request_id,
message: match result {
Ok(message) => Ok(message),
Err(tokio::time::Elapsed { .. }) => {
debug!(
"[{}] Response did not complete before deadline of {}s.",
self.ctx.trace_id(),
format_rfc3339(self.deadline)
);
// No point in responding, since the client will have dropped the
// request.
Err(ServerError {
kind: io::ErrorKind::TimedOut,
detail: Some(format!(
"Response did not complete before deadline of {}s.",
format_rfc3339(self.deadline)
)),
_non_exhaustive: (),
})
}
},
_non_exhaustive: (),
});
*self.as_mut().project().state = RespState::PollReady;
}
RespState::PollReady => {
let ready = ready!(self.as_mut().project().response_tx.poll_ready(cx));
if ready.is_err() {
return Poll::Ready(());
}
let resp = (self.ctx, self.as_mut().project().response.take().unwrap());
if self
.as_mut()
.project()
.response_tx
.start_send(resp)
.is_err()
{
return Poll::Ready(());
}
*self.as_mut().project().state = RespState::PollFlush;
}
RespState::PollFlush => {
let ready = ready!(self.as_mut().project().response_tx.poll_flush(cx));
if ready.is_err() {
return Poll::Ready(());
}
return Poll::Ready(());
}
}
}
}
}
impl<C, S> Stream for ClientHandler<C, S>
where
C: Channel,
S: Serve<C::Req, Resp = C::Resp>,
{
type Item = io::Result<RequestHandler<S::Fut, C::Resp>>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
let read = self.as_mut().pump_read(cx)?;
let read_closed = if let Poll::Ready(None) = read {
true
} else {
false
};
match (read, self.as_mut().pump_write(cx, read_closed)?) {
(Poll::Ready(None), Poll::Ready(None)) => {
return Poll::Ready(None);
}
(Poll::Ready(Some(request_handler)), _) => {
return Poll::Ready(Some(Ok(request_handler)));
}
(_, Poll::Ready(Some(()))) => {}
_ => {
return Poll::Pending;
}
}
}
}
}
// Send + 'static execution helper methods.
impl<C, S> ClientHandler<C, S>
where
C: Channel + 'static,
C::Req: Send + 'static,
C::Resp: Send + 'static,
S: Serve<C::Req, Resp = C::Resp> + Send + 'static,
S::Fut: Send + 'static,
{
/// Runs the client handler until completion by spawning each
/// request handler onto the default executor.
#[cfg(feature = "tokio1")]
pub fn execute(self) -> impl Future<Output = ()> {
use log::info;
self.try_for_each(|request_handler| {
async {
tokio::spawn(request_handler);
Ok(())
}
})
.unwrap_or_else(|e| info!("ClientHandler errored out: {}", e))
}
}
/// A future that drives the server by spawning channels and request handlers on the default
/// executor.
#[pin_project]
#[derive(Debug)]
#[cfg(feature = "tokio1")]
pub struct Running<St, Se> {
#[pin]
incoming: St,
server: Se,
}
#[cfg(feature = "tokio1")]
impl<St, C, Se> Future for Running<St, Se>
where
St: Sized + Stream<Item = C>,
C: Channel + Send + 'static,
C::Req: Send + 'static,
C::Resp: Send + 'static,
Se: Serve<C::Req, Resp = C::Resp> + Send + 'static + Clone,
Se::Fut: Send + 'static,
{
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
use log::info;
while let Some(channel) = ready!(self.as_mut().project().incoming.poll_next(cx)) {
tokio::spawn(
channel
.respond_with(self.as_mut().project().server.clone())
.execute(),
);
}
info!("Server shutting down.");
Poll::Ready(())
}
}

View File

@@ -1,326 +0,0 @@
use super::{Channel, Config};
use crate::{Response, ServerError};
use futures::{future::AbortRegistration, prelude::*, ready, task::*};
use log::debug;
use pin_project::pin_project;
use std::{io, pin::Pin};
/// A [`Channel`] that limits the number of concurrent
/// requests by throttling.
#[pin_project]
#[derive(Debug)]
pub struct Throttler<C> {
max_in_flight_requests: usize,
#[pin]
inner: C,
}
impl<C> Throttler<C> {
/// Returns the inner channel.
pub fn get_ref(&self) -> &C {
&self.inner
}
}
impl<C> Throttler<C>
where
C: Channel,
{
/// Returns a new `Throttler` that wraps the given channel and limits concurrent requests to
/// `max_in_flight_requests`.
pub fn new(inner: C, max_in_flight_requests: usize) -> Self {
Throttler {
inner,
max_in_flight_requests,
}
}
}
impl<C> Stream for Throttler<C>
where
C: Channel,
{
type Item = <C as Stream>::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
while self.as_mut().in_flight_requests() >= *self.as_mut().project().max_in_flight_requests
{
ready!(self.as_mut().project().inner.poll_ready(cx)?);
match ready!(self.as_mut().project().inner.poll_next(cx)?) {
Some(request) => {
debug!(
"[{}] Client has reached in-flight request limit ({}/{}).",
request.context.trace_id(),
self.as_mut().in_flight_requests(),
self.as_mut().project().max_in_flight_requests,
);
self.as_mut().start_send(Response {
request_id: request.id,
message: Err(ServerError {
kind: io::ErrorKind::WouldBlock,
detail: Some("Server throttled the request.".into()),
_non_exhaustive: (),
}),
_non_exhaustive: (),
})?;
}
None => return Poll::Ready(None),
}
}
self.project().inner.poll_next(cx)
}
}
impl<C> Sink<Response<<C as Channel>::Resp>> for Throttler<C>
where
C: Channel,
{
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().inner.poll_ready(cx)
}
fn start_send(self: Pin<&mut Self>, item: Response<<C as Channel>::Resp>) -> io::Result<()> {
self.project().inner.start_send(item)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
self.project().inner.poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<io::Result<()>> {
self.project().inner.poll_close(cx)
}
}
impl<C> AsRef<C> for Throttler<C> {
fn as_ref(&self) -> &C {
&self.inner
}
}
impl<C> Channel for Throttler<C>
where
C: Channel,
{
type Req = <C as Channel>::Req;
type Resp = <C as Channel>::Resp;
fn in_flight_requests(self: Pin<&mut Self>) -> usize {
self.project().inner.in_flight_requests()
}
fn config(&self) -> &Config {
self.inner.config()
}
fn start_request(self: Pin<&mut Self>, request_id: u64) -> AbortRegistration {
self.project().inner.start_request(request_id)
}
}
/// A stream of throttling channels.
#[pin_project]
#[derive(Debug)]
pub struct ThrottlerStream<S> {
#[pin]
inner: S,
max_in_flight_requests: usize,
}
impl<S> ThrottlerStream<S>
where
S: Stream,
<S as Stream>::Item: Channel,
{
pub(crate) fn new(inner: S, max_in_flight_requests: usize) -> Self {
Self {
inner,
max_in_flight_requests,
}
}
}
impl<S> Stream for ThrottlerStream<S>
where
S: Stream,
<S as Stream>::Item: Channel,
{
type Item = Throttler<<S as Stream>::Item>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
match ready!(self.as_mut().project().inner.poll_next(cx)) {
Some(channel) => Poll::Ready(Some(Throttler::new(
channel,
*self.project().max_in_flight_requests,
))),
None => Poll::Ready(None),
}
}
}
#[cfg(test)]
use super::testing::{self, FakeChannel, PollExt};
#[cfg(test)]
use crate::Request;
#[cfg(test)]
use pin_utils::pin_mut;
#[cfg(test)]
use std::marker::PhantomData;
#[test]
fn throttler_in_flight_requests() {
let throttler = Throttler {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
for i in 0..5 {
throttler.inner.in_flight_requests.insert(i);
}
assert_eq!(throttler.as_mut().in_flight_requests(), 5);
}
#[test]
fn throttler_start_request() {
let throttler = Throttler {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler.as_mut().start_request(1);
assert_eq!(throttler.inner.in_flight_requests.len(), 1);
}
#[test]
fn throttler_poll_next_done() {
let throttler = Throttler {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
assert!(throttler.as_mut().poll_next(&mut testing::cx()).is_done());
}
#[test]
fn throttler_poll_next_some() -> io::Result<()> {
let throttler = Throttler {
max_in_flight_requests: 1,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler.inner.push_req(0, 1);
assert!(throttler.as_mut().poll_ready(&mut testing::cx()).is_ready());
assert_eq!(
throttler
.as_mut()
.poll_next(&mut testing::cx())?
.map(|r| r.map(|r| (r.id, r.message))),
Poll::Ready(Some((0, 1)))
);
Ok(())
}
#[test]
fn throttler_poll_next_throttled() {
let throttler = Throttler {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler.inner.push_req(1, 1);
assert!(throttler.as_mut().poll_next(&mut testing::cx()).is_done());
assert_eq!(throttler.inner.sink.len(), 1);
let resp = throttler.inner.sink.get(0).unwrap();
assert_eq!(resp.request_id, 1);
assert!(resp.message.is_err());
}
#[test]
fn throttler_poll_next_throttled_sink_not_ready() {
let throttler = Throttler {
max_in_flight_requests: 0,
inner: PendingSink::default::<isize, isize>(),
};
pin_mut!(throttler);
assert!(throttler.poll_next(&mut testing::cx()).is_pending());
struct PendingSink<In, Out> {
ghost: PhantomData<fn(Out) -> In>,
}
impl PendingSink<(), ()> {
pub fn default<Req, Resp>() -> PendingSink<io::Result<Request<Req>>, Response<Resp>> {
PendingSink { ghost: PhantomData }
}
}
impl<In, Out> Stream for PendingSink<In, Out> {
type Item = In;
fn poll_next(self: Pin<&mut Self>, _: &mut Context) -> Poll<Option<Self::Item>> {
unimplemented!()
}
}
impl<In, Out> Sink<Out> for PendingSink<In, Out> {
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
fn start_send(self: Pin<&mut Self>, _: Out) -> Result<(), Self::Error> {
Err(io::Error::from(io::ErrorKind::WouldBlock))
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
fn poll_close(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
}
impl<Req, Resp> Channel for PendingSink<io::Result<Request<Req>>, Response<Resp>> {
type Req = Req;
type Resp = Resp;
fn config(&self) -> &Config {
unimplemented!()
}
fn in_flight_requests(self: Pin<&mut Self>) -> usize {
0
}
fn start_request(self: Pin<&mut Self>, _: u64) -> AbortRegistration {
unimplemented!()
}
}
}
#[test]
fn throttler_start_send() {
let throttler = Throttler {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler.inner.in_flight_requests.insert(0);
throttler
.as_mut()
.start_send(Response {
request_id: 0,
message: Ok(1),
_non_exhaustive: (),
})
.unwrap();
assert!(throttler.inner.in_flight_requests.is_empty());
assert_eq!(
throttler.inner.sink.get(0),
Some(&Response {
request_id: 0,
message: Ok(1),
_non_exhaustive: ()
})
);
}

View File

@@ -1,123 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Transports backed by in-memory channels.
use crate::PollIo;
use futures::{channel::mpsc, task::*, Sink, Stream};
use pin_project::pin_project;
use std::io;
use std::pin::Pin;
/// Returns two unbounded channel peers. Each [`Stream`] yields items sent through the other's
/// [`Sink`].
pub fn unbounded<SinkItem, Item>() -> (
UnboundedChannel<SinkItem, Item>,
UnboundedChannel<Item, SinkItem>,
) {
let (tx1, rx2) = mpsc::unbounded();
let (tx2, rx1) = mpsc::unbounded();
(
UnboundedChannel { tx: tx1, rx: rx1 },
UnboundedChannel { tx: tx2, rx: rx2 },
)
}
/// A bi-directional channel backed by an [`UnboundedSender`](mpsc::UnboundedSender)
/// and [`UnboundedReceiver`](mpsc::UnboundedReceiver).
#[pin_project]
#[derive(Debug)]
pub struct UnboundedChannel<Item, SinkItem> {
#[pin]
rx: mpsc::UnboundedReceiver<Item>,
#[pin]
tx: mpsc::UnboundedSender<SinkItem>,
}
impl<Item, SinkItem> Stream for UnboundedChannel<Item, SinkItem> {
type Item = Result<Item, io::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> PollIo<Item> {
self.project().rx.poll_next(cx).map(|option| option.map(Ok))
}
}
impl<Item, SinkItem> Sink<SinkItem> for UnboundedChannel<Item, SinkItem> {
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project()
.tx
.poll_ready(cx)
.map_err(|_| io::Error::from(io::ErrorKind::NotConnected))
}
fn start_send(self: Pin<&mut Self>, item: SinkItem) -> io::Result<()> {
self.project()
.tx
.start_send(item)
.map_err(|_| io::Error::from(io::ErrorKind::NotConnected))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project()
.tx
.poll_flush(cx)
.map_err(|_| io::Error::from(io::ErrorKind::NotConnected))
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project()
.tx
.poll_close(cx)
.map_err(|_| io::Error::from(io::ErrorKind::NotConnected))
}
}
#[cfg(test)]
mod tests {
use crate::{
client, context,
server::{Handler, Server},
transport,
};
use assert_matches::assert_matches;
use futures::{prelude::*, stream};
use log::trace;
use std::io;
#[cfg(feature = "tokio1")]
#[tokio::test(threaded_scheduler)]
async fn integration() -> io::Result<()> {
let _ = env_logger::try_init();
let (client_channel, server_channel) = transport::channel::unbounded();
tokio::spawn(
Server::default()
.incoming(stream::once(future::ready(server_channel)))
.respond_with(|_ctx, request: String| {
future::ready(request.parse::<u64>().map_err(|_| {
io::Error::new(
io::ErrorKind::InvalidInput,
format!("{:?} is not an int", request),
)
}))
}),
);
let mut client = client::new(client::Config::default(), client_channel).spawn()?;
let response1 = client.call(context::current(), "123".into()).await?;
let response2 = client.call(context::current(), "abc".into()).await?;
trace!("response1: {:?}, response2: {:?}", response1, response2);
assert_matches!(response1, Ok(123));
assert_matches!(response2, Err(ref e) if e.kind() == io::ErrorKind::InvalidInput);
Ok(())
}
}

View File

@@ -1,30 +0,0 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a [`Transport`] trait as well as implementations.
//!
//! The rpc crate is transport- and protocol-agnostic. Any transport that impls [`Transport`]
//! can be plugged in, using whatever protocol it wants.
use futures::prelude::*;
use std::io;
pub mod channel;
pub(crate) mod sealed {
use super::*;
/// A bidirectional stream ([`Sink`] + [`Stream`]) of messages.
pub trait Transport<SinkItem, Item>:
Stream<Item = io::Result<Item>> + Sink<SinkItem, Error = io::Error>
{
}
impl<T, SinkItem, Item> Transport<SinkItem, Item> for T where
T: Stream<Item = io::Result<Item>> + Sink<SinkItem, Error = io::Error> + ?Sized
{
}
}

View File

@@ -0,0 +1,735 @@
// Copyright 2019 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! A generic Serde-based `Transport` that can serialize anything supported by `tokio-serde` via any medium that implements `AsyncRead` and `AsyncWrite`.
#![deny(missing_docs)]
use futures::{prelude::*, task::*};
use pin_project::pin_project;
use serde::{Deserialize, Serialize};
use std::{error::Error, io, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_serde::{Framed as SerdeFramed, *};
use tokio_util::codec::{length_delimited::LengthDelimitedCodec, Framed};
/// A transport that serializes to, and deserializes from, a byte stream.
#[pin_project]
pub struct Transport<S, Item, SinkItem, Codec> {
#[pin]
inner: SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>,
}
impl<S, Item, SinkItem, Codec> Transport<S, Item, SinkItem, Codec> {
/// Returns the inner transport over which messages are sent and received.
pub fn get_ref(&self) -> &S {
self.inner.get_ref().get_ref()
}
}
impl<S, Item, SinkItem, Codec, CodecError> Stream for Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite + AsyncRead,
Item: for<'a> Deserialize<'a>,
Codec: Deserializer<Item>,
CodecError: Into<Box<dyn std::error::Error + Send + Sync>>,
SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
Stream<Item = Result<Item, CodecError>>,
{
type Item = io::Result<Item>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<io::Result<Item>>> {
self.project()
.inner
.poll_next(cx)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
}
impl<S, Item, SinkItem, Codec, CodecError> Sink<SinkItem> for Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite,
SinkItem: Serialize,
Codec: Serializer<SinkItem>,
CodecError: Into<Box<dyn Error + Send + Sync>>,
SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
Sink<SinkItem, Error = CodecError>,
{
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project()
.inner
.poll_ready(cx)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
fn start_send(self: Pin<&mut Self>, item: SinkItem) -> io::Result<()> {
self.project()
.inner
.start_send(item)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project()
.inner
.poll_flush(cx)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
self.project()
.inner
.poll_close(cx)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
}
/// Constructs a new transport from a framed transport and a serialization codec.
pub fn new<S, Item, SinkItem, Codec>(
framed_io: Framed<S, LengthDelimitedCodec>,
codec: Codec,
) -> Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite + AsyncRead,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
{
Transport {
inner: SerdeFramed::new(framed_io, codec),
}
}
impl<S, Item, SinkItem, Codec> From<(S, Codec)> for Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite + AsyncRead,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
{
fn from((io, codec): (S, Codec)) -> Self {
new(Framed::new(io, LengthDelimitedCodec::new()), codec)
}
}
#[cfg(feature = "tcp")]
#[cfg_attr(docsrs, doc(cfg(feature = "tcp")))]
/// TCP support for generic transport using Tokio.
pub mod tcp {
use {
super::*,
futures::ready,
std::{marker::PhantomData, net::SocketAddr},
tokio::net::{TcpListener, TcpStream, ToSocketAddrs},
tokio_util::codec::length_delimited,
};
impl<Item, SinkItem, Codec> Transport<TcpStream, Item, SinkItem, Codec> {
/// Returns the peer address of the underlying TcpStream.
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.inner.get_ref().get_ref().peer_addr()
}
/// Returns the local address of the underlying TcpStream.
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.inner.get_ref().get_ref().local_addr()
}
}
/// A connection Future that also exposes the length-delimited framing config.
#[must_use]
#[pin_project]
pub struct Connect<T, Item, SinkItem, CodecFn> {
#[pin]
inner: T,
codec_fn: CodecFn,
config: length_delimited::Builder,
ghost: PhantomData<(fn(SinkItem), fn() -> Item)>,
}
impl<T, Item, SinkItem, Codec, CodecFn> Future for Connect<T, Item, SinkItem, CodecFn>
where
T: Future<Output = io::Result<TcpStream>>,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
type Output = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let io = ready!(self.as_mut().project().inner.poll(cx))?;
Poll::Ready(Ok(new(self.config.new_framed(io), (self.codec_fn)())))
}
}
impl<T, Item, SinkItem, CodecFn> Connect<T, Item, SinkItem, CodecFn> {
/// Returns an immutable reference to the length-delimited codec's config.
pub fn config(&self) -> &length_delimited::Builder {
&self.config
}
/// Returns a mutable reference to the length-delimited codec's config.
pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
&mut self.config
}
}
/// Connects to `addr`, wrapping the connection in a TCP transport.
pub fn connect<A, Item, SinkItem, Codec, CodecFn>(
addr: A,
codec_fn: CodecFn,
) -> Connect<impl Future<Output = io::Result<TcpStream>>, Item, SinkItem, CodecFn>
where
A: ToSocketAddrs,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
Connect {
inner: TcpStream::connect(addr),
codec_fn,
config: LengthDelimitedCodec::builder(),
ghost: PhantomData,
}
}
/// Listens on `addr`, wrapping accepted connections in TCP transports.
pub async fn listen<A, Item, SinkItem, Codec, CodecFn>(
addr: A,
codec_fn: CodecFn,
) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
where
A: ToSocketAddrs,
Item: for<'de> Deserialize<'de>,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
listen_on(TcpListener::bind(addr).await?, codec_fn).await
}
/// Wrap accepted connections from `listener` in TCP transports.
pub async fn listen_on<Item, SinkItem, Codec, CodecFn>(
listener: TcpListener,
codec_fn: CodecFn,
) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
where
Item: for<'de> Deserialize<'de>,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
let local_addr = listener.local_addr()?;
Ok(Incoming {
listener,
codec_fn,
local_addr,
config: LengthDelimitedCodec::builder(),
ghost: PhantomData,
})
}
/// A [`TcpListener`] that wraps connections in [transports](Transport).
#[pin_project]
#[derive(Debug)]
pub struct Incoming<Item, SinkItem, Codec, CodecFn> {
listener: TcpListener,
local_addr: SocketAddr,
codec_fn: CodecFn,
config: length_delimited::Builder,
ghost: PhantomData<(fn() -> Item, fn(SinkItem), Codec)>,
}
impl<Item, SinkItem, Codec, CodecFn> Incoming<Item, SinkItem, Codec, CodecFn> {
/// Returns the address being listened on.
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
/// Returns an immutable reference to the length-delimited codec's config.
pub fn config(&self) -> &length_delimited::Builder {
&self.config
}
/// Returns a mutable reference to the length-delimited codec's config.
pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
&mut self.config
}
}
impl<Item, SinkItem, Codec, CodecFn> Stream for Incoming<Item, SinkItem, Codec, CodecFn>
where
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
type Item = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let conn: TcpStream =
ready!(Pin::new(&mut self.as_mut().project().listener).poll_accept(cx)?).0;
Poll::Ready(Some(Ok(new(
self.config.new_framed(conn),
(self.codec_fn)(),
))))
}
}
}
#[cfg(all(unix, feature = "unix"))]
#[cfg_attr(docsrs, doc(cfg(all(unix, feature = "unix"))))]
/// Unix Domain Socket support for generic transport using Tokio.
pub mod unix {
use {
super::*,
futures::ready,
std::{marker::PhantomData, path::Path},
tokio::net::{unix::SocketAddr, UnixListener, UnixStream},
tokio_util::codec::length_delimited,
};
impl<Item, SinkItem, Codec> Transport<UnixStream, Item, SinkItem, Codec> {
/// Returns the socket address of the remote half of the underlying [`UnixStream`].
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.inner.get_ref().get_ref().peer_addr()
}
/// Returns the socket address of the local half of the underlying [`UnixStream`].
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.inner.get_ref().get_ref().local_addr()
}
}
/// A connection Future that also exposes the length-delimited framing config.
#[must_use]
#[pin_project]
pub struct Connect<T, Item, SinkItem, CodecFn> {
#[pin]
inner: T,
codec_fn: CodecFn,
config: length_delimited::Builder,
ghost: PhantomData<(fn(SinkItem), fn() -> Item)>,
}
impl<T, Item, SinkItem, Codec, CodecFn> Future for Connect<T, Item, SinkItem, CodecFn>
where
T: Future<Output = io::Result<UnixStream>>,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
type Output = io::Result<Transport<UnixStream, Item, SinkItem, Codec>>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
let io = ready!(self.as_mut().project().inner.poll(cx))?;
Poll::Ready(Ok(new(self.config.new_framed(io), (self.codec_fn)())))
}
}
impl<T, Item, SinkItem, CodecFn> Connect<T, Item, SinkItem, CodecFn> {
/// Returns an immutable reference to the length-delimited codec's config.
pub fn config(&self) -> &length_delimited::Builder {
&self.config
}
/// Returns a mutable reference to the length-delimited codec's config.
pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
&mut self.config
}
}
/// Connects to socket named by `path`, wrapping the connection in a Unix Domain Socket
/// transport.
pub fn connect<P, Item, SinkItem, Codec, CodecFn>(
path: P,
codec_fn: CodecFn,
) -> Connect<impl Future<Output = io::Result<UnixStream>>, Item, SinkItem, CodecFn>
where
P: AsRef<Path>,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
Connect {
inner: UnixStream::connect(path),
codec_fn,
config: LengthDelimitedCodec::builder(),
ghost: PhantomData,
}
}
/// Listens on the socket named by `path`, wrapping accepted connections in Unix Domain Socket
/// transports.
pub async fn listen<P, Item, SinkItem, Codec, CodecFn>(
path: P,
codec_fn: CodecFn,
) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
where
P: AsRef<Path>,
Item: for<'de> Deserialize<'de>,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
listen_on(UnixListener::bind(path)?, codec_fn).await
}
/// Wrap accepted connections from `listener` in Unix Domain Socket transports.
pub async fn listen_on<Item, SinkItem, Codec, CodecFn>(
listener: UnixListener,
codec_fn: CodecFn,
) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
where
Item: for<'de> Deserialize<'de>,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
let local_addr = listener.local_addr()?;
Ok(Incoming {
listener,
codec_fn,
local_addr,
config: LengthDelimitedCodec::builder(),
ghost: PhantomData,
})
}
/// A [`UnixListener`] that wraps connections in [transports](Transport).
#[pin_project]
#[derive(Debug)]
pub struct Incoming<Item, SinkItem, Codec, CodecFn> {
listener: UnixListener,
local_addr: SocketAddr,
codec_fn: CodecFn,
config: length_delimited::Builder,
ghost: PhantomData<(fn() -> Item, fn(SinkItem), Codec)>,
}
impl<Item, SinkItem, Codec, CodecFn> Incoming<Item, SinkItem, Codec, CodecFn> {
/// Returns the the socket address being listened on.
pub fn local_addr(&self) -> &SocketAddr {
&self.local_addr
}
/// Returns an immutable reference to the length-delimited codec's config.
pub fn config(&self) -> &length_delimited::Builder {
&self.config
}
/// Returns a mutable reference to the length-delimited codec's config.
pub fn config_mut(&mut self) -> &mut length_delimited::Builder {
&mut self.config
}
}
impl<Item, SinkItem, Codec, CodecFn> Stream for Incoming<Item, SinkItem, Codec, CodecFn>
where
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
type Item = io::Result<Transport<UnixStream, Item, SinkItem, Codec>>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let conn: UnixStream = ready!(self.as_mut().project().listener.poll_accept(cx)?).0;
Poll::Ready(Some(Ok(new(
self.config.new_framed(conn),
(self.codec_fn)(),
))))
}
}
/// A temporary `PathBuf` that lives in `std::env::temp_dir` and is removed on drop.
pub struct TempPathBuf(std::path::PathBuf);
impl TempPathBuf {
/// A named socket that results in `<tempdir>/<name>`
pub fn new<S: AsRef<str>>(name: S) -> Self {
let mut sock = std::env::temp_dir();
sock.push(name.as_ref());
Self(sock)
}
/// Appends a random hex string to the socket name resulting in
/// `<tempdir>/<name>_<xxxxx>`
pub fn with_random<S: AsRef<str>>(name: S) -> Self {
Self::new(format!("{}_{:016x}", name.as_ref(), rand::random::<u64>()))
}
}
impl AsRef<std::path::Path> for TempPathBuf {
fn as_ref(&self) -> &std::path::Path {
self.0.as_path()
}
}
impl Drop for TempPathBuf {
fn drop(&mut self) {
// This will remove the file pointed to by this PathBuf if it exists, however Err's can
// be returned such as attempting to remove a non-existing file, or one which we don't
// have permission to remove. In these cases the Err is swallowed
let _ = std::fs::remove_file(&self.0);
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio_serde::formats::SymmetricalJson;
#[test]
fn temp_path_buf_non_random() {
let sock = TempPathBuf::new("test");
let mut good = std::env::temp_dir();
good.push("test");
assert_eq!(sock.as_ref(), good);
assert_eq!(sock.as_ref().file_name().unwrap(), "test");
}
#[test]
fn temp_path_buf_random() {
let sock = TempPathBuf::with_random("test");
let good = std::env::temp_dir();
assert!(sock.as_ref().starts_with(good));
// Since there are 16 random characters we just assert the file_name has the right name
// and starts with the correct string 'test_'
// file name: test_xxxxxxxxxxxxxxxx
// test = 4
// _ = 1
// <hex> = 16
// total = 21
let fname = sock.as_ref().file_name().unwrap().to_string_lossy();
assert!(fname.starts_with("test_"));
assert_eq!(fname.len(), 21);
}
#[test]
fn temp_path_buf_non_existing() {
let sock = TempPathBuf::with_random("test");
let sock_path = std::path::PathBuf::from(sock.as_ref());
// No actual file has been created yet
assert!(!sock_path.exists());
// Should not panic
std::mem::drop(sock);
assert!(!sock_path.exists());
}
#[test]
fn temp_path_buf_existing_file() {
let sock = TempPathBuf::with_random("test");
let sock_path = std::path::PathBuf::from(sock.as_ref());
let _file = std::fs::File::create(&sock).unwrap();
assert!(sock_path.exists());
std::mem::drop(sock);
assert!(!sock_path.exists());
}
#[test]
fn temp_path_buf_preexisting_file() {
let mut pre_existing = std::env::temp_dir();
pre_existing.push("test");
let _file = std::fs::File::create(&pre_existing).unwrap();
let sock = TempPathBuf::new("test");
let sock_path = std::path::PathBuf::from(sock.as_ref());
assert!(sock_path.exists());
std::mem::drop(sock);
assert!(!sock_path.exists());
}
#[tokio::test]
async fn temp_path_buf_for_socket() {
let sock = TempPathBuf::with_random("test");
// Save path for testing after drop
let sock_path = std::path::PathBuf::from(sock.as_ref());
// create the actual socket
let _ = listen(&sock, SymmetricalJson::<String>::default).await;
assert!(sock_path.exists());
std::mem::drop(sock);
assert!(!sock_path.exists());
}
}
}
#[cfg(test)]
mod tests {
use super::Transport;
use assert_matches::assert_matches;
use futures::{task::*, Sink, SinkExt, Stream, StreamExt};
use pin_utils::pin_mut;
use std::{
io::{self, Cursor},
pin::Pin,
};
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
use tokio_serde::formats::SymmetricalJson;
fn ctx() -> Context<'static> {
Context::from_waker(noop_waker_ref())
}
struct TestIo(Cursor<Vec<u8>>);
impl AsyncRead for TestIo {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
AsyncRead::poll_read(Pin::new(&mut self.0), cx, buf)
}
}
impl AsyncWrite for TestIo {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
AsyncWrite::poll_write(Pin::new(&mut self.0), cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
AsyncWrite::poll_flush(Pin::new(&mut self.0), cx)
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
AsyncWrite::poll_shutdown(Pin::new(&mut self.0), cx)
}
}
#[test]
fn close() {
let (tx, _rx) = crate::transport::channel::bounded::<(), ()>(0);
pin_mut!(tx);
assert_matches!(tx.as_mut().poll_close(&mut ctx()), Poll::Ready(Ok(())));
assert_matches!(tx.as_mut().start_send(()), Err(_));
}
#[test]
fn test_stream() {
let data: &[u8] = b"\x00\x00\x00\x18\"Test one, check check.\"";
let transport = Transport::from((
TestIo(Cursor::new(Vec::from(data))),
SymmetricalJson::<String>::default(),
));
pin_mut!(transport);
assert_matches!(
transport.as_mut().poll_next(&mut ctx()),
Poll::Ready(Some(Ok(ref s))) if s == "Test one, check check.");
assert_matches!(transport.as_mut().poll_next(&mut ctx()), Poll::Ready(None));
}
#[test]
fn test_sink() {
let writer = Cursor::new(vec![]);
let mut transport = Box::pin(Transport::from((
TestIo(writer),
SymmetricalJson::<String>::default(),
)));
assert_matches!(
transport.as_mut().poll_ready(&mut ctx()),
Poll::Ready(Ok(()))
);
assert_matches!(
transport
.as_mut()
.start_send("Test one, check check.".into()),
Ok(())
);
assert_matches!(
transport.as_mut().poll_flush(&mut ctx()),
Poll::Ready(Ok(()))
);
assert_eq!(
transport.get_ref().0.get_ref(),
b"\x00\x00\x00\x18\"Test one, check check.\""
);
}
#[cfg(feature = "tcp")]
#[tokio::test]
async fn tcp() -> io::Result<()> {
use super::tcp;
let mut listener = tcp::listen("0.0.0.0:0", SymmetricalJson::<String>::default).await?;
let addr = listener.local_addr();
tokio::spawn(async move {
let mut transport = listener.next().await.unwrap().unwrap();
let message = transport.next().await.unwrap().unwrap();
transport.send(message).await.unwrap();
});
let mut transport = tcp::connect(addr, SymmetricalJson::<String>::default).await?;
transport.send(String::from("test")).await?;
assert_matches!(transport.next().await, Some(Ok(s)) if s == "test");
assert_matches!(transport.next().await, None);
Ok(())
}
#[cfg(feature = "tcp")]
#[tokio::test]
async fn tcp_on_existing_transport() -> io::Result<()> {
use super::tcp;
let transport = tokio::net::TcpListener::bind("0.0.0.0:0").await?;
let mut listener = tcp::listen_on(transport, SymmetricalJson::<String>::default).await?;
let addr = listener.local_addr();
tokio::spawn(async move {
let mut transport = listener.next().await.unwrap().unwrap();
let message = transport.next().await.unwrap().unwrap();
transport.send(message).await.unwrap();
});
let mut transport = tcp::connect(addr, SymmetricalJson::<String>::default).await?;
transport.send(String::from("test")).await?;
assert_matches!(transport.next().await, Some(Ok(s)) if s == "test");
assert_matches!(transport.next().await, None);
Ok(())
}
#[cfg(all(unix, feature = "unix"))]
#[tokio::test]
async fn uds() -> io::Result<()> {
use super::unix;
let sock = unix::TempPathBuf::with_random("uds");
let mut listener = unix::listen(&sock, SymmetricalJson::<String>::default).await?;
tokio::spawn(async move {
let mut transport = listener.next().await.unwrap().unwrap();
let message = transport.next().await.unwrap().unwrap();
transport.send(message).await.unwrap();
});
let mut transport = unix::connect(&sock, SymmetricalJson::<String>::default).await?;
transport.send(String::from("test")).await?;
assert_matches!(transport.next().await, Some(Ok(s)) if s == "test");
assert_matches!(transport.next().await, None);
Ok(())
}
#[cfg(all(unix, feature = "unix"))]
#[tokio::test]
async fn uds_on_existing_transport() -> io::Result<()> {
use super::unix;
let sock = unix::TempPathBuf::with_random("uds");
let transport = tokio::net::UnixListener::bind(&sock)?;
let mut listener = unix::listen_on(transport, SymmetricalJson::<String>::default).await?;
tokio::spawn(async move {
let mut transport = listener.next().await.unwrap().unwrap();
let message = transport.next().await.unwrap().unwrap();
transport.send(message).await.unwrap();
});
let mut transport = unix::connect(&sock, SymmetricalJson::<String>::default).await?;
transport.send(String::from("test")).await?;
assert_matches!(transport.next().await, Some(Ok(s)) if s == "test");
assert_matches!(transport.next().await, None);
Ok(())
}
}

View File

@@ -1,326 +0,0 @@
// Copyright 2019 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! A generic Serde-based `Transport` that can serialize anything supported by `tokio-serde` via any medium that implements `AsyncRead` and `AsyncWrite`.
#![deny(missing_docs)]
use futures::{prelude::*, task::*};
use pin_project::pin_project;
use serde::{Deserialize, Serialize};
use std::{error::Error, io, pin::Pin};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_serde::{Framed as SerdeFramed, *};
use tokio_util::codec::{length_delimited::LengthDelimitedCodec, Framed};
/// A transport that serializes to, and deserializes from, a [`TcpStream`].
#[pin_project]
pub struct Transport<S, Item, SinkItem, Codec> {
#[pin]
inner: SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>,
}
impl<S, Item, SinkItem, Codec, CodecError> Stream for Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite + AsyncRead,
Item: for<'a> Deserialize<'a>,
Codec: Deserializer<Item>,
CodecError: Into<Box<dyn std::error::Error + Send + Sync>>,
SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
Stream<Item = Result<Item, CodecError>>,
{
type Item = io::Result<Item>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<io::Result<Item>>> {
match self.project().inner.poll_next(cx) {
Poll::Pending => Poll::Pending,
Poll::Ready(None) => Poll::Ready(None),
Poll::Ready(Some(Ok::<_, CodecError>(next))) => Poll::Ready(Some(Ok(next))),
Poll::Ready(Some(Err::<_, CodecError>(e))) => {
Poll::Ready(Some(Err(io::Error::new(io::ErrorKind::Other, e))))
}
}
}
}
impl<S, Item, SinkItem, Codec, CodecError> Sink<SinkItem> for Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite,
SinkItem: Serialize,
Codec: Serializer<SinkItem>,
CodecError: Into<Box<dyn Error + Send + Sync>>,
SerdeFramed<Framed<S, LengthDelimitedCodec>, Item, SinkItem, Codec>:
Sink<SinkItem, Error = CodecError>,
{
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
convert(self.project().inner.poll_ready(cx))
}
fn start_send(self: Pin<&mut Self>, item: SinkItem) -> io::Result<()> {
self.project()
.inner
.start_send(item)
.map_err(|e| io::Error::new(io::ErrorKind::Other, e))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
convert(self.project().inner.poll_flush(cx))
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
convert(self.project().inner.poll_close(cx))
}
}
fn convert<E: Into<Box<dyn Error + Send + Sync>>>(
poll: Poll<Result<(), E>>,
) -> Poll<io::Result<()>> {
poll.map(|ready| ready.map_err(|e| io::Error::new(io::ErrorKind::Other, e)))
}
impl<S, Item, SinkItem, Codec> From<(S, Codec)> for Transport<S, Item, SinkItem, Codec>
where
S: AsyncWrite + AsyncRead,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
{
fn from((inner, codec): (S, Codec)) -> Self {
Transport {
inner: SerdeFramed::new(Framed::new(inner, LengthDelimitedCodec::new()), codec),
}
}
}
#[cfg(feature = "tcp")]
#[cfg_attr(docsrs, doc(cfg(feature = "tcp")))]
/// TCP support for generic transport using Tokio.
pub mod tcp {
use {
super::*,
futures::ready,
std::{marker::PhantomData, net::SocketAddr},
tokio::net::{TcpListener, TcpStream, ToSocketAddrs},
};
mod private {
use super::*;
pub trait Sealed {}
impl<Item, SinkItem, Codec> Sealed for Transport<TcpStream, Item, SinkItem, Codec> {}
}
impl<Item, SinkItem, Codec> Transport<TcpStream, Item, SinkItem, Codec> {
/// Returns the peer address of the underlying TcpStream.
pub fn peer_addr(&self) -> io::Result<SocketAddr> {
self.inner.get_ref().get_ref().peer_addr()
}
/// Returns the local address of the underlying TcpStream.
pub fn local_addr(&self) -> io::Result<SocketAddr> {
self.inner.get_ref().get_ref().local_addr()
}
}
/// Returns a new JSON transport that reads from and writes to `io`.
pub fn new<Item, SinkItem, Codec>(
io: TcpStream,
codec: Codec,
) -> Transport<TcpStream, Item, SinkItem, Codec>
where
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
{
Transport::from((io, codec))
}
/// Connects to `addr`, wrapping the connection in a JSON transport.
pub async fn connect<A, Item, SinkItem, Codec>(
addr: A,
codec: Codec,
) -> io::Result<Transport<TcpStream, Item, SinkItem, Codec>>
where
A: ToSocketAddrs,
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
{
Ok(new(TcpStream::connect(addr).await?, codec))
}
/// Listens on `addr`, wrapping accepted connections in JSON transports.
pub async fn listen<A, Item, SinkItem, Codec, CodecFn>(
addr: A,
codec_fn: CodecFn,
) -> io::Result<Incoming<Item, SinkItem, Codec, CodecFn>>
where
A: ToSocketAddrs,
Item: for<'de> Deserialize<'de>,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
let listener = TcpListener::bind(addr).await?;
let local_addr = listener.local_addr()?;
Ok(Incoming {
listener,
codec_fn,
local_addr,
ghost: PhantomData,
})
}
/// A [`TcpListener`] that wraps connections in JSON transports.
#[pin_project]
#[derive(Debug)]
pub struct Incoming<Item, SinkItem, Codec, CodecFn> {
listener: TcpListener,
local_addr: SocketAddr,
codec_fn: CodecFn,
ghost: PhantomData<(Item, SinkItem, Codec)>,
}
impl<Item, SinkItem, Codec, CodecFn> Incoming<Item, SinkItem, Codec, CodecFn> {
/// Returns the address being listened on.
pub fn local_addr(&self) -> SocketAddr {
self.local_addr
}
}
impl<Item, SinkItem, Codec, CodecFn> Stream for Incoming<Item, SinkItem, Codec, CodecFn>
where
Item: for<'de> Deserialize<'de>,
SinkItem: Serialize,
Codec: Serializer<SinkItem> + Deserializer<Item>,
CodecFn: Fn() -> Codec,
{
type Item = io::Result<Transport<TcpStream, Item, SinkItem, Codec>>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let next =
ready!(Pin::new(&mut self.as_mut().project().listener.incoming()).poll_next(cx)?);
Poll::Ready(next.map(|conn| Ok(new(conn, (self.codec_fn)()))))
}
}
}
#[cfg(test)]
mod tests {
use super::Transport;
use assert_matches::assert_matches;
use futures::{task::*, Sink, Stream};
use pin_utils::pin_mut;
use std::{
io::{self, Cursor},
pin::Pin,
};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_serde::formats::SymmetricalJson;
fn ctx() -> Context<'static> {
Context::from_waker(&noop_waker_ref())
}
#[test]
fn test_stream() {
struct TestIo(Cursor<&'static [u8]>);
impl AsyncRead for TestIo {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &mut [u8],
) -> Poll<io::Result<usize>> {
AsyncRead::poll_read(Pin::new(self.0.get_mut()), cx, buf)
}
}
impl AsyncWrite for TestIo {
fn poll_write(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_buf: &[u8],
) -> Poll<io::Result<usize>> {
unreachable!()
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
unreachable!()
}
fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
unreachable!()
}
}
let data = b"\x00\x00\x00\x18\"Test one, check check.\"";
let transport = Transport::from((
TestIo(Cursor::new(data)),
SymmetricalJson::<String>::default(),
));
pin_mut!(transport);
assert_matches!(
transport.poll_next(&mut ctx()),
Poll::Ready(Some(Ok(ref s))) if s == "Test one, check check.");
}
#[test]
fn test_sink() {
struct TestIo<'a>(&'a mut Vec<u8>);
impl<'a> AsyncRead for TestIo<'a> {
fn poll_read(
self: Pin<&mut Self>,
_cx: &mut Context<'_>,
_buf: &mut [u8],
) -> Poll<io::Result<usize>> {
unreachable!()
}
}
impl<'a> AsyncWrite for TestIo<'a> {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<io::Result<usize>> {
AsyncWrite::poll_write(Pin::new(&mut *self.0), cx, buf)
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
AsyncWrite::poll_flush(Pin::new(&mut *self.0), cx)
}
fn poll_shutdown(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<io::Result<()>> {
AsyncWrite::poll_shutdown(Pin::new(&mut *self.0), cx)
}
}
let mut writer = vec![];
let transport =
Transport::from((TestIo(&mut writer), SymmetricalJson::<String>::default()));
pin_mut!(transport);
assert_matches!(
transport.as_mut().poll_ready(&mut ctx()),
Poll::Ready(Ok(()))
);
assert_matches!(
transport
.as_mut()
.start_send("Test one, check check.".into()),
Ok(())
);
assert_matches!(transport.poll_flush(&mut ctx()), Poll::Ready(Ok(())));
assert_eq!(writer, b"\x00\x00\x00\x18\"Test one, check check.\"");
}
}

1592
tarpc/src/server.rs Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,221 @@
use crate::util::{Compact, TimeUntil};
use fnv::FnvHashMap;
use futures::future::{AbortHandle, AbortRegistration};
use std::{
collections::hash_map,
task::{Context, Poll},
time::SystemTime,
};
use tokio_util::time::delay_queue::{self, DelayQueue};
use tracing::Span;
/// A data structure that tracks in-flight requests. It aborts requests,
/// either on demand or when a request deadline expires.
#[derive(Debug, Default)]
pub struct InFlightRequests {
request_data: FnvHashMap<u64, RequestData>,
deadlines: DelayQueue<u64>,
}
/// Data needed to clean up a single in-flight request.
#[derive(Debug)]
struct RequestData {
/// Aborts the response handler for the associated request.
abort_handle: AbortHandle,
/// The key to remove the timer for the request's deadline.
deadline_key: delay_queue::Key,
/// The client span.
span: Span,
}
/// An error returned when a request attempted to start with the same ID as a request already
/// in flight.
#[derive(Debug)]
pub struct AlreadyExistsError;
impl InFlightRequests {
/// Returns the number of in-flight requests.
pub fn len(&self) -> usize {
self.request_data.len()
}
/// Starts a request, unless a request with the same ID is already in flight.
pub fn start_request(
&mut self,
request_id: u64,
deadline: SystemTime,
span: Span,
) -> Result<AbortRegistration, AlreadyExistsError> {
match self.request_data.entry(request_id) {
hash_map::Entry::Vacant(vacant) => {
let timeout = deadline.time_until();
let (abort_handle, abort_registration) = AbortHandle::new_pair();
let deadline_key = self.deadlines.insert(request_id, timeout);
vacant.insert(RequestData {
abort_handle,
deadline_key,
span,
});
Ok(abort_registration)
}
hash_map::Entry::Occupied(_) => Err(AlreadyExistsError),
}
}
/// Cancels an in-flight request. Returns true iff the request was found.
pub fn cancel_request(&mut self, request_id: u64) -> bool {
if let Some(RequestData {
span,
abort_handle,
deadline_key,
}) = self.request_data.remove(&request_id)
{
let _entered = span.enter();
self.request_data.compact(0.1);
abort_handle.abort();
self.deadlines.remove(&deadline_key);
tracing::info!("ReceiveCancel");
true
} else {
false
}
}
/// Removes a request without aborting. Returns true iff the request was found.
/// This method should be used when a response is being sent.
pub fn remove_request(&mut self, request_id: u64) -> Option<Span> {
if let Some(request_data) = self.request_data.remove(&request_id) {
self.request_data.compact(0.1);
self.deadlines.remove(&request_data.deadline_key);
Some(request_data.span)
} else {
None
}
}
/// Yields a request that has expired, aborting any ongoing processing of that request.
pub fn poll_expired(&mut self, cx: &mut Context) -> Poll<Option<u64>> {
if self.deadlines.is_empty() {
// TODO(https://github.com/tokio-rs/tokio/issues/4161)
// This is a workaround for DelayQueue not always treating this case correctly.
return Poll::Ready(None);
}
self.deadlines.poll_expired(cx).map(|expired| {
let expired = expired?;
if let Some(RequestData {
abort_handle, span, ..
}) = self.request_data.remove(expired.get_ref())
{
let _entered = span.enter();
self.request_data.compact(0.1);
abort_handle.abort();
tracing::error!("DeadlineExceeded");
}
Some(expired.into_inner())
})
}
}
/// When InFlightRequests is dropped, any outstanding requests are aborted.
impl Drop for InFlightRequests {
fn drop(&mut self) {
self.request_data
.values()
.for_each(|request_data| request_data.abort_handle.abort())
}
}
#[cfg(test)]
mod tests {
use super::*;
use assert_matches::assert_matches;
use futures::{
future::{pending, Abortable},
FutureExt,
};
use futures_test::task::noop_context;
#[tokio::test]
async fn start_request_increases_len() {
let mut in_flight_requests = InFlightRequests::default();
assert_eq!(in_flight_requests.len(), 0);
in_flight_requests
.start_request(0, SystemTime::now(), Span::current())
.unwrap();
assert_eq!(in_flight_requests.len(), 1);
}
#[tokio::test]
async fn polling_expired_aborts() {
let mut in_flight_requests = InFlightRequests::default();
let abort_registration = in_flight_requests
.start_request(0, SystemTime::now(), Span::current())
.unwrap();
let mut abortable_future = Box::new(Abortable::new(pending::<()>(), abort_registration));
tokio::time::pause();
tokio::time::advance(std::time::Duration::from_secs(1000)).await;
assert_matches!(
in_flight_requests.poll_expired(&mut noop_context()),
Poll::Ready(Some(_))
);
assert_matches!(
abortable_future.poll_unpin(&mut noop_context()),
Poll::Ready(Err(_))
);
assert_eq!(in_flight_requests.len(), 0);
}
#[tokio::test]
async fn cancel_request_aborts() {
let mut in_flight_requests = InFlightRequests::default();
let abort_registration = in_flight_requests
.start_request(0, SystemTime::now(), Span::current())
.unwrap();
let mut abortable_future = Box::new(Abortable::new(pending::<()>(), abort_registration));
assert!(in_flight_requests.cancel_request(0));
assert_matches!(
abortable_future.poll_unpin(&mut noop_context()),
Poll::Ready(Err(_))
);
assert_eq!(in_flight_requests.len(), 0);
}
#[tokio::test]
async fn remove_request_doesnt_abort() {
let mut in_flight_requests = InFlightRequests::default();
assert!(in_flight_requests.deadlines.is_empty());
let abort_registration = in_flight_requests
.start_request(
0,
SystemTime::now() + std::time::Duration::from_secs(10),
Span::current(),
)
.unwrap();
let mut abortable_future = Box::new(Abortable::new(pending::<()>(), abort_registration));
// Precondition: Pending expiration
assert_matches!(
in_flight_requests.poll_expired(&mut noop_context()),
Poll::Pending
);
assert!(!in_flight_requests.deadlines.is_empty());
assert_matches!(in_flight_requests.remove_request(0), Some(_));
// Postcondition: No pending expirations
assert!(in_flight_requests.deadlines.is_empty());
assert_matches!(
in_flight_requests.poll_expired(&mut noop_context()),
Poll::Ready(None)
);
assert_matches!(
abortable_future.poll_unpin(&mut noop_context()),
Poll::Pending
);
assert_eq!(in_flight_requests.len(), 0);
}
}

View File

@@ -0,0 +1,92 @@
use super::{
limits::{channels_per_key::MaxChannelsPerKey, requests_per_channel::MaxRequestsPerChannel},
Channel, Serve,
};
use futures::prelude::*;
use std::{fmt, hash::Hash};
/// An extension trait for [streams](futures::prelude::Stream) of [`Channels`](Channel).
pub trait Incoming<C>
where
Self: Sized + Stream<Item = C>,
C: Channel,
{
/// Enforces channel per-key limits.
fn max_channels_per_key<K, KF>(self, n: u32, keymaker: KF) -> MaxChannelsPerKey<Self, K, KF>
where
K: fmt::Display + Eq + Hash + Clone + Unpin,
KF: Fn(&C) -> K,
{
MaxChannelsPerKey::new(self, n, keymaker)
}
/// Caps the number of concurrent requests per channel.
fn max_concurrent_requests_per_channel(self, n: usize) -> MaxRequestsPerChannel<Self> {
MaxRequestsPerChannel::new(self, n)
}
/// Returns a stream of channels in execution. Each channel in execution is a stream of
/// futures, where each future is an in-flight request being rsponded to.
fn execute<S>(
self,
serve: S,
) -> impl Stream<Item = impl Stream<Item = impl Future<Output = ()>>>
where
S: Serve<Req = C::Req, Resp = C::Resp> + Clone,
{
self.map(move |channel| channel.execute(serve.clone()))
}
}
#[cfg(feature = "tokio1")]
/// Spawns all channels-in-execution, delegating to the tokio runtime to manage their completion.
/// Each channel is spawned, and each request from each channel is spawned.
/// Note that this function is generic over any stream-of-streams-of-futures, but it is intended
/// for spawning streams of channels.
///
/// # Example
/// ```rust
/// use tarpc::{
/// context,
/// client::{self, NewClient},
/// server::{self, BaseChannel, Channel, incoming::{Incoming, spawn_incoming}, serve},
/// transport,
/// };
/// use futures::prelude::*;
///
/// #[tokio::main]
/// async fn main() {
/// let (tx, rx) = transport::channel::unbounded();
/// let NewClient { client, dispatch } = client::new(client::Config::default(), tx);
/// tokio::spawn(dispatch);
///
/// let incoming = stream::once(async move {
/// BaseChannel::new(server::Config::default(), rx)
/// }).execute(serve(|_, i| async move { Ok(i + 1) }));
/// tokio::spawn(spawn_incoming(incoming));
/// assert_eq!(client.call(context::current(), "AddOne", 1).await.unwrap(), 2);
/// }
/// ```
pub async fn spawn_incoming(
incoming: impl Stream<
Item = impl Stream<Item = impl Future<Output = ()> + Send + 'static> + Send + 'static,
>,
) {
use futures::pin_mut;
pin_mut!(incoming);
while let Some(channel) = incoming.next().await {
tokio::spawn(async move {
pin_mut!(channel);
while let Some(request) = channel.next().await {
tokio::spawn(request);
}
});
}
}
impl<S, C> Incoming<C> for S
where
S: Sized + Stream<Item = C>,
C: Channel,
{
}

View File

@@ -0,0 +1,5 @@
/// Provides functionality to limit the number of active channels.
pub mod channels_per_key;
/// Provides a [channel](crate::server::Channel) that limits the number of in-flight requests.
pub mod requests_per_channel;

View File

@@ -9,78 +9,65 @@ use crate::{
util::Compact,
};
use fnv::FnvHashMap;
use futures::{channel::mpsc, future::AbortRegistration, prelude::*, ready, stream::Fuse, task::*};
use log::{debug, info, trace};
use futures::{prelude::*, ready, stream::Fuse, task::*};
use pin_project::pin_project;
use raii_counter::{Counter, WeakCounter};
use std::sync::{Arc, Weak};
use std::{
collections::hash_map::Entry, convert::TryInto, fmt, hash::Hash, marker::Unpin, pin::Pin,
collections::hash_map::Entry, convert::TryFrom, fmt, hash::Hash, marker::Unpin, pin::Pin,
};
use tokio::sync::mpsc;
use tracing::{debug, info, trace};
/// A single-threaded filter that drops channels based on per-key limits.
/// An [`Incoming`](crate::server::incoming::Incoming) stream that drops new channels based on
/// per-key limits.
///
/// The decision to drop a Channel is made once at the time the Channel materializes. Once a
/// Channel is yielded, it will not be prematurely dropped.
#[pin_project]
#[derive(Debug)]
pub struct ChannelFilter<S, K, F>
pub struct MaxChannelsPerKey<S, K, F>
where
K: Eq + Hash,
{
#[pin]
listener: Fuse<S>,
channels_per_key: u32,
#[pin]
dropped_keys: mpsc::UnboundedReceiver<K>,
#[pin]
dropped_keys_tx: mpsc::UnboundedSender<K>,
key_counts: FnvHashMap<K, TrackerPrototype<K>>,
key_counts: FnvHashMap<K, Weak<Tracker<K>>>,
keymaker: F,
}
/// A channel that is tracked by a ChannelFilter.
/// A channel that is tracked by [`MaxChannelsPerKey`].
#[pin_project]
#[derive(Debug)]
pub struct TrackedChannel<C, K> {
#[pin]
inner: C,
tracker: Tracker<K>,
tracker: Arc<Tracker<K>>,
}
#[derive(Clone, Debug)]
#[derive(Debug)]
struct Tracker<K> {
key: Option<Arc<K>>,
counter: Counter,
key: Option<K>,
dropped_keys: mpsc::UnboundedSender<K>,
}
impl<K> Drop for Tracker<K> {
fn drop(&mut self) {
if self.counter.count() <= 1 {
// Don't care if the listener is dropped.
match Arc::try_unwrap(self.key.take().unwrap()) {
Ok(key) => {
let _ = self.dropped_keys.unbounded_send(key);
}
_ => unreachable!(),
}
}
// Don't care if the listener is dropped.
let _ = self.dropped_keys.send(self.key.take().unwrap());
}
}
#[derive(Clone, Debug)]
struct TrackerPrototype<K> {
key: Weak<K>,
counter: WeakCounter,
dropped_keys: mpsc::UnboundedSender<K>,
}
impl<C, K> Stream for TrackedChannel<C, K>
where
C: Stream,
{
type Item = <C as Stream>::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
self.channel().poll_next(cx)
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
self.inner_pin_mut().poll_next(cx)
}
}
@@ -90,20 +77,20 @@ where
{
type Error = C::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.channel().poll_ready(cx)
fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.inner_pin_mut().poll_ready(cx)
}
fn start_send(self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
self.channel().start_send(item)
fn start_send(mut self: Pin<&mut Self>, item: I) -> Result<(), Self::Error> {
self.inner_pin_mut().start_send(item)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.channel().poll_flush(cx)
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.inner_pin_mut().poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.channel().poll_close(cx)
fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.inner_pin_mut().poll_close(cx)
}
}
@@ -119,17 +106,18 @@ where
{
type Req = C::Req;
type Resp = C::Resp;
type Transport = C::Transport;
fn config(&self) -> &server::Config {
self.inner.config()
}
fn in_flight_requests(self: Pin<&mut Self>) -> usize {
self.project().inner.in_flight_requests()
fn in_flight_requests(&self) -> usize {
self.inner.in_flight_requests()
}
fn start_request(self: Pin<&mut Self>, request_id: u64) -> AbortRegistration {
self.project().inner.start_request(request_id)
fn transport(&self) -> &Self::Transport {
self.inner.transport()
}
}
@@ -140,12 +128,12 @@ impl<C, K> TrackedChannel<C, K> {
}
/// Returns the pinned inner channel.
fn channel<'a>(self: Pin<&'a mut Self>) -> Pin<&'a mut C> {
self.project().inner
fn inner_pin_mut<'a>(self: &'a mut Pin<&mut Self>) -> Pin<&'a mut C> {
self.as_mut().project().inner
}
}
impl<S, K, F> ChannelFilter<S, K, F>
impl<S, K, F> MaxChannelsPerKey<S, K, F>
where
K: Eq + Hash,
S: Stream,
@@ -153,8 +141,8 @@ where
{
/// Sheds new channels to stay under configured limits.
pub(crate) fn new(listener: S, channels_per_key: u32, keymaker: F) -> Self {
let (dropped_keys_tx, dropped_keys) = mpsc::unbounded();
ChannelFilter {
let (dropped_keys_tx, dropped_keys) = mpsc::unbounded_channel();
MaxChannelsPerKey {
listener: listener.fuse(),
channels_per_key,
dropped_keys,
@@ -165,12 +153,16 @@ where
}
}
impl<S, K, F> ChannelFilter<S, K, F>
impl<S, K, F> MaxChannelsPerKey<S, K, F>
where
S: Stream,
K: fmt::Display + Eq + Hash + Clone + Unpin,
F: Fn(&S::Item) -> K,
{
fn listener_pin_mut<'a>(self: &'a mut Pin<&mut Self>) -> Pin<&'a mut Fuse<S>> {
self.as_mut().project().listener
}
fn handle_new_channel(
mut self: Pin<&mut Self>,
stream: S::Item,
@@ -179,11 +171,10 @@ where
let tracker = self.as_mut().increment_channels_for_key(key.clone())?;
trace!(
"[{}] Opening channel ({}/{}) channels for key.",
key,
tracker.counter.count(),
self.as_mut().project().channels_per_key
);
channel_filter_key = %key,
open_channels = Arc::strong_count(&tracker),
max_open_channels = self.channels_per_key,
"Opening channel");
Ok(TrackedChannel {
tracker,
@@ -191,45 +182,38 @@ where
})
}
fn increment_channels_for_key(mut self: Pin<&mut Self>, key: K) -> Result<Tracker<K>, K> {
let channels_per_key = self.channels_per_key;
let dropped_keys = self.dropped_keys_tx.clone();
let key_counts = &mut self.as_mut().project().key_counts;
match key_counts.entry(key.clone()) {
fn increment_channels_for_key(self: Pin<&mut Self>, key: K) -> Result<Arc<Tracker<K>>, K> {
let self_ = self.project();
let dropped_keys = self_.dropped_keys_tx;
match self_.key_counts.entry(key.clone()) {
Entry::Vacant(vacant) => {
let key = Arc::new(key);
let counter = WeakCounter::new();
vacant.insert(TrackerPrototype {
key: Arc::downgrade(&key),
counter: counter.clone(),
let tracker = Arc::new(Tracker {
key: Some(key),
dropped_keys: dropped_keys.clone(),
});
Ok(Tracker {
key: Some(key),
counter: counter.upgrade(),
dropped_keys,
})
vacant.insert(Arc::downgrade(&tracker));
Ok(tracker)
}
Entry::Occupied(o) => {
let count = o.get().counter.count();
if count >= channels_per_key.try_into().unwrap() {
Entry::Occupied(mut o) => {
let count = o.get().strong_count();
if count >= TryFrom::try_from(*self_.channels_per_key).unwrap() {
info!(
"[{}] Opened max channels from key ({}/{}).",
key, count, channels_per_key
);
channel_filter_key = %key,
open_channels = count,
max_open_channels = *self_.channels_per_key,
"At open channel limit");
Err(key)
} else {
let TrackerPrototype {
key,
counter,
dropped_keys,
} = o.get().clone();
Ok(Tracker {
counter: counter.upgrade(),
key: Some(key.upgrade().unwrap()),
dropped_keys,
})
Ok(o.get().upgrade().unwrap_or_else(|| {
let tracker = Arc::new(Tracker {
key: Some(key),
dropped_keys: dropped_keys.clone(),
});
*o.get_mut() = Arc::downgrade(&tracker);
tracker
}))
}
}
}
@@ -239,18 +223,21 @@ where
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<TrackedChannel<S::Item, K>, K>>> {
match ready!(self.as_mut().project().listener.poll_next_unpin(cx)) {
match ready!(self.listener_pin_mut().poll_next_unpin(cx)) {
Some(codec) => Poll::Ready(Some(self.handle_new_channel(codec))),
None => Poll::Ready(None),
}
}
fn poll_closed_channels(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
match ready!(self.as_mut().project().dropped_keys.poll_next_unpin(cx)) {
fn poll_closed_channels(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let self_ = self.project();
match ready!(self_.dropped_keys.poll_recv(cx)) {
Some(key) => {
debug!("All channels dropped for key [{}]", key);
self.as_mut().project().key_counts.remove(&key);
self.as_mut().project().key_counts.compact(0.1);
debug!(
channel_filter_key = %key,
"All channels dropped");
self_.key_counts.remove(&key);
self_.key_counts.compact(0.1);
Poll::Ready(())
}
None => unreachable!("Holding a copy of closed_channels and didn't close it."),
@@ -258,7 +245,7 @@ where
}
}
impl<S, K, F> Stream for ChannelFilter<S, K, F>
impl<S, K, F> Stream for MaxChannelsPerKey<S, K, F>
where
S: Stream,
K: fmt::Display + Eq + Hash + Clone + Unpin,
@@ -291,43 +278,38 @@ where
}
}
}
#[cfg(test)]
fn ctx() -> Context<'static> {
use futures::task::*;
Context::from_waker(&noop_waker_ref())
Context::from_waker(noop_waker_ref())
}
#[test]
fn tracker_drop() {
use assert_matches::assert_matches;
use raii_counter::Counter;
let (tx, mut rx) = mpsc::unbounded();
let (tx, mut rx) = mpsc::unbounded_channel();
Tracker {
key: Some(Arc::new(1)),
counter: Counter::new(),
key: Some(1),
dropped_keys: tx,
};
assert_matches!(rx.try_next(), Ok(Some(1)));
assert_matches!(rx.poll_recv(&mut ctx()), Poll::Ready(Some(1)));
}
#[test]
fn tracked_channel_stream() {
use assert_matches::assert_matches;
use pin_utils::pin_mut;
use raii_counter::Counter;
let (chan_tx, chan) = mpsc::unbounded();
let (dropped_keys, _) = mpsc::unbounded();
let (chan_tx, chan) = futures::channel::mpsc::unbounded();
let (dropped_keys, _) = mpsc::unbounded_channel();
let channel = TrackedChannel {
inner: chan,
tracker: Tracker {
key: Some(Arc::new(1)),
counter: Counter::new(),
tracker: Arc::new(Tracker {
key: Some(1),
dropped_keys,
},
}),
};
chan_tx.unbounded_send("test").unwrap();
@@ -339,17 +321,15 @@ fn tracked_channel_stream() {
fn tracked_channel_sink() {
use assert_matches::assert_matches;
use pin_utils::pin_mut;
use raii_counter::Counter;
let (chan, mut chan_rx) = mpsc::unbounded();
let (dropped_keys, _) = mpsc::unbounded();
let (chan, mut chan_rx) = futures::channel::mpsc::unbounded();
let (dropped_keys, _) = mpsc::unbounded_channel();
let channel = TrackedChannel {
inner: chan,
tracker: Tracker {
key: Some(Arc::new(1)),
counter: Counter::new(),
tracker: Arc::new(Tracker {
key: Some(1),
dropped_keys,
},
}),
};
pin_mut!(channel);
@@ -367,16 +347,16 @@ fn channel_filter_increment_channels_for_key() {
struct TestChannel {
key: &'static str,
}
let (_, listener) = mpsc::unbounded();
let filter = ChannelFilter::new(listener, 2, |chan: &TestChannel| chan.key);
let (_, listener) = futures::channel::mpsc::unbounded();
let filter = MaxChannelsPerKey::new(listener, 2, |chan: &TestChannel| chan.key);
pin_mut!(filter);
let tracker1 = filter.as_mut().increment_channels_for_key("key").unwrap();
assert_eq!(tracker1.counter.count(), 1);
assert_eq!(Arc::strong_count(&tracker1), 1);
let tracker2 = filter.as_mut().increment_channels_for_key("key").unwrap();
assert_eq!(tracker1.counter.count(), 2);
assert_eq!(Arc::strong_count(&tracker1), 2);
assert_matches!(filter.increment_channels_for_key("key"), Err("key"));
drop(tracker2);
assert_eq!(tracker1.counter.count(), 1);
assert_eq!(Arc::strong_count(&tracker1), 1);
}
#[test]
@@ -388,27 +368,27 @@ fn channel_filter_handle_new_channel() {
struct TestChannel {
key: &'static str,
}
let (_, listener) = mpsc::unbounded();
let filter = ChannelFilter::new(listener, 2, |chan: &TestChannel| chan.key);
let (_, listener) = futures::channel::mpsc::unbounded();
let filter = MaxChannelsPerKey::new(listener, 2, |chan: &TestChannel| chan.key);
pin_mut!(filter);
let channel1 = filter
.as_mut()
.handle_new_channel(TestChannel { key: "key" })
.unwrap();
assert_eq!(channel1.tracker.counter.count(), 1);
assert_eq!(Arc::strong_count(&channel1.tracker), 1);
let channel2 = filter
.as_mut()
.handle_new_channel(TestChannel { key: "key" })
.unwrap();
assert_eq!(channel1.tracker.counter.count(), 2);
assert_eq!(Arc::strong_count(&channel1.tracker), 2);
assert_matches!(
filter.handle_new_channel(TestChannel { key: "key" }),
Err("key")
);
drop(channel2);
assert_eq!(channel1.tracker.counter.count(), 1);
assert_eq!(Arc::strong_count(&channel1.tracker), 1);
}
#[test]
@@ -420,8 +400,8 @@ fn channel_filter_poll_listener() {
struct TestChannel {
key: &'static str,
}
let (new_channels, listener) = mpsc::unbounded();
let filter = ChannelFilter::new(listener, 2, |chan: &TestChannel| chan.key);
let (new_channels, listener) = futures::channel::mpsc::unbounded();
let filter = MaxChannelsPerKey::new(listener, 2, |chan: &TestChannel| chan.key);
pin_mut!(filter);
new_channels
@@ -429,14 +409,14 @@ fn channel_filter_poll_listener() {
.unwrap();
let channel1 =
assert_matches!(filter.as_mut().poll_listener(&mut ctx()), Poll::Ready(Some(Ok(c))) => c);
assert_eq!(channel1.tracker.counter.count(), 1);
assert_eq!(Arc::strong_count(&channel1.tracker), 1);
new_channels
.unbounded_send(TestChannel { key: "key" })
.unwrap();
let _channel2 =
assert_matches!(filter.as_mut().poll_listener(&mut ctx()), Poll::Ready(Some(Ok(c))) => c);
assert_eq!(channel1.tracker.counter.count(), 2);
assert_eq!(Arc::strong_count(&channel1.tracker), 2);
new_channels
.unbounded_send(TestChannel { key: "key" })
@@ -444,7 +424,7 @@ fn channel_filter_poll_listener() {
let key =
assert_matches!(filter.as_mut().poll_listener(&mut ctx()), Poll::Ready(Some(Err(k))) => k);
assert_eq!(key, "key");
assert_eq!(channel1.tracker.counter.count(), 2);
assert_eq!(Arc::strong_count(&channel1.tracker), 2);
}
#[test]
@@ -456,8 +436,8 @@ fn channel_filter_poll_closed_channels() {
struct TestChannel {
key: &'static str,
}
let (new_channels, listener) = mpsc::unbounded();
let filter = ChannelFilter::new(listener, 2, |chan: &TestChannel| chan.key);
let (new_channels, listener) = futures::channel::mpsc::unbounded();
let filter = MaxChannelsPerKey::new(listener, 2, |chan: &TestChannel| chan.key);
pin_mut!(filter);
new_channels
@@ -484,8 +464,8 @@ fn channel_filter_stream() {
struct TestChannel {
key: &'static str,
}
let (new_channels, listener) = mpsc::unbounded();
let filter = ChannelFilter::new(listener, 2, |chan: &TestChannel| chan.key);
let (new_channels, listener) = futures::channel::mpsc::unbounded();
let filter = MaxChannelsPerKey::new(listener, 2, |chan: &TestChannel| chan.key);
pin_mut!(filter);
new_channels

View File

@@ -0,0 +1,349 @@
// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use crate::{
server::{Channel, Config},
Response, ServerError,
};
use futures::{prelude::*, ready, task::*};
use pin_project::pin_project;
use std::{io, pin::Pin};
/// A [`Channel`] that limits the number of concurrent requests by throttling.
///
/// Note that this is a very basic throttling heuristic. It is easy to set a number that is too low
/// for the resources available to the server. For production use cases, a more advanced throttler
/// is likely needed.
#[pin_project]
#[derive(Debug)]
pub struct MaxRequests<C> {
max_in_flight_requests: usize,
#[pin]
inner: C,
}
impl<C> MaxRequests<C> {
/// Returns the inner channel.
pub fn get_ref(&self) -> &C {
&self.inner
}
}
impl<C> MaxRequests<C>
where
C: Channel,
{
/// Returns a new `MaxRequests` that wraps the given channel and limits concurrent requests to
/// `max_in_flight_requests`.
pub fn new(inner: C, max_in_flight_requests: usize) -> Self {
MaxRequests {
max_in_flight_requests,
inner,
}
}
}
impl<C> Stream for MaxRequests<C>
where
C: Channel,
{
type Item = <C as Stream>::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
while self.as_mut().in_flight_requests() >= *self.as_mut().project().max_in_flight_requests
{
ready!(self.as_mut().project().inner.poll_ready(cx)?);
match ready!(self.as_mut().project().inner.poll_next(cx)?) {
Some(r) => {
let _entered = r.span.enter();
tracing::info!(
in_flight_requests = self.as_mut().in_flight_requests(),
"ThrottleRequest",
);
self.as_mut().start_send(Response {
request_id: r.request.id,
message: Err(ServerError {
kind: io::ErrorKind::WouldBlock,
detail: "server throttled the request.".into(),
}),
})?;
}
None => return Poll::Ready(None),
}
}
self.project().inner.poll_next(cx)
}
}
impl<C> Sink<Response<<C as Channel>::Resp>> for MaxRequests<C>
where
C: Channel,
{
type Error = C::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().inner.poll_ready(cx)
}
fn start_send(
self: Pin<&mut Self>,
item: Response<<C as Channel>::Resp>,
) -> Result<(), Self::Error> {
self.project().inner.start_send(item)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().inner.poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Result<(), Self::Error>> {
self.project().inner.poll_close(cx)
}
}
impl<C> AsRef<C> for MaxRequests<C> {
fn as_ref(&self) -> &C {
&self.inner
}
}
impl<C> Channel for MaxRequests<C>
where
C: Channel,
{
type Req = <C as Channel>::Req;
type Resp = <C as Channel>::Resp;
type Transport = <C as Channel>::Transport;
fn in_flight_requests(&self) -> usize {
self.inner.in_flight_requests()
}
fn config(&self) -> &Config {
self.inner.config()
}
fn transport(&self) -> &Self::Transport {
self.inner.transport()
}
}
/// An [`Incoming`](crate::server::incoming::Incoming) stream of channels that enforce limits on
/// the number of in-flight requests.
#[pin_project]
#[derive(Debug)]
pub struct MaxRequestsPerChannel<S> {
#[pin]
inner: S,
max_in_flight_requests: usize,
}
impl<S> MaxRequestsPerChannel<S>
where
S: Stream,
<S as Stream>::Item: Channel,
{
pub(crate) fn new(inner: S, max_in_flight_requests: usize) -> Self {
Self {
inner,
max_in_flight_requests,
}
}
}
impl<S> Stream for MaxRequestsPerChannel<S>
where
S: Stream,
<S as Stream>::Item: Channel,
{
type Item = MaxRequests<<S as Stream>::Item>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
match ready!(self.as_mut().project().inner.poll_next(cx)) {
Some(channel) => Poll::Ready(Some(MaxRequests::new(
channel,
*self.project().max_in_flight_requests,
))),
None => Poll::Ready(None),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::server::{
testing::{self, FakeChannel, PollExt},
TrackedRequest,
};
use pin_utils::pin_mut;
use std::{
marker::PhantomData,
time::{Duration, SystemTime},
};
use tracing::Span;
#[tokio::test]
async fn throttler_in_flight_requests() {
let throttler = MaxRequests {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
for i in 0..5 {
throttler
.inner
.in_flight_requests
.start_request(
i,
SystemTime::now() + Duration::from_secs(1),
Span::current(),
)
.unwrap();
}
assert_eq!(throttler.as_mut().in_flight_requests(), 5);
}
#[test]
fn throttler_poll_next_done() {
let throttler = MaxRequests {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
assert!(throttler.as_mut().poll_next(&mut testing::cx()).is_done());
}
#[test]
fn throttler_poll_next_some() -> io::Result<()> {
let throttler = MaxRequests {
max_in_flight_requests: 1,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler.inner.push_req(0, 1);
assert!(throttler.as_mut().poll_ready(&mut testing::cx()).is_ready());
assert_eq!(
throttler
.as_mut()
.poll_next(&mut testing::cx())?
.map(|r| r.map(|r| (r.request.id, r.request.message))),
Poll::Ready(Some((0, 1)))
);
Ok(())
}
#[test]
fn throttler_poll_next_throttled() {
let throttler = MaxRequests {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler.inner.push_req(1, 1);
assert!(throttler.as_mut().poll_next(&mut testing::cx()).is_done());
assert_eq!(throttler.inner.sink.len(), 1);
let resp = throttler.inner.sink.get(0).unwrap();
assert_eq!(resp.request_id, 1);
assert!(resp.message.is_err());
}
#[test]
fn throttler_poll_next_throttled_sink_not_ready() {
let throttler = MaxRequests {
max_in_flight_requests: 0,
inner: PendingSink::default::<isize, isize>(),
};
pin_mut!(throttler);
assert!(throttler.poll_next(&mut testing::cx()).is_pending());
struct PendingSink<In, Out> {
ghost: PhantomData<fn(Out) -> In>,
}
impl PendingSink<(), ()> {
pub fn default<Req, Resp>(
) -> PendingSink<io::Result<TrackedRequest<Req>>, Response<Resp>> {
PendingSink { ghost: PhantomData }
}
}
impl<In, Out> Stream for PendingSink<In, Out> {
type Item = In;
fn poll_next(self: Pin<&mut Self>, _: &mut Context) -> Poll<Option<Self::Item>> {
unimplemented!()
}
}
impl<In, Out> Sink<Out> for PendingSink<In, Out> {
type Error = io::Error;
fn poll_ready(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
fn start_send(self: Pin<&mut Self>, _: Out) -> Result<(), Self::Error> {
Err(io::Error::from(io::ErrorKind::WouldBlock))
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
fn poll_close(self: Pin<&mut Self>, _: &mut Context) -> Poll<Result<(), Self::Error>> {
Poll::Pending
}
}
impl<Req, Resp> Channel for PendingSink<io::Result<TrackedRequest<Req>>, Response<Resp>> {
type Req = Req;
type Resp = Resp;
type Transport = ();
fn config(&self) -> &Config {
unimplemented!()
}
fn in_flight_requests(&self) -> usize {
0
}
fn transport(&self) -> &() {
&()
}
}
}
#[tokio::test]
async fn throttler_start_send() {
let throttler = MaxRequests {
max_in_flight_requests: 0,
inner: FakeChannel::default::<isize, isize>(),
};
pin_mut!(throttler);
throttler
.inner
.in_flight_requests
.start_request(
0,
SystemTime::now() + Duration::from_secs(1),
Span::current(),
)
.unwrap();
throttler
.as_mut()
.start_send(Response {
request_id: 0,
message: Ok(1),
})
.unwrap();
assert_eq!(throttler.inner.in_flight_requests.len(), 0);
assert_eq!(
throttler.inner.sink.get(0),
Some(&Response {
request_id: 0,
message: Ok(1),
})
);
}
}

View File

@@ -0,0 +1,25 @@
// Copyright 2022 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Hooks for horizontal functionality that can run either before or after a request is executed.
/// A request hook that runs before a request is executed.
mod before;
/// A request hook that runs after a request is completed.
mod after;
/// A request hook that runs both before a request is executed and after it is completed.
mod before_and_after;
pub use {
after::{AfterRequest, ServeThenHook},
before::{
before, BeforeRequest, BeforeRequestCons, BeforeRequestList, BeforeRequestNil,
HookThenServe,
},
before_and_after::HookThenServeThenHook,
};

View File

@@ -0,0 +1,72 @@
// Copyright 2022 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a hook that runs after request execution.
use crate::{context, server::Serve, ServerError};
use futures::prelude::*;
/// A hook that runs after request execution.
#[allow(async_fn_in_trait)]
pub trait AfterRequest<Resp> {
/// The function that is called after request execution.
///
/// The hook can modify the request context and the response.
async fn after(&mut self, ctx: &mut context::Context, resp: &mut Result<Resp, ServerError>);
}
impl<F, Fut, Resp> AfterRequest<Resp> for F
where
F: FnMut(&mut context::Context, &mut Result<Resp, ServerError>) -> Fut,
Fut: Future<Output = ()>,
{
async fn after(&mut self, ctx: &mut context::Context, resp: &mut Result<Resp, ServerError>) {
self(ctx, resp).await
}
}
/// A Service function that runs a hook after request execution.
pub struct ServeThenHook<Serv, Hook> {
serve: Serv,
hook: Hook,
}
impl<Serv, Hook> ServeThenHook<Serv, Hook> {
pub(crate) fn new(serve: Serv, hook: Hook) -> Self {
Self { serve, hook }
}
}
impl<Serv: Clone, Hook: Clone> Clone for ServeThenHook<Serv, Hook> {
fn clone(&self) -> Self {
Self {
serve: self.serve.clone(),
hook: self.hook.clone(),
}
}
}
impl<Serv, Hook> Serve for ServeThenHook<Serv, Hook>
where
Serv: Serve,
Hook: AfterRequest<Serv::Resp>,
{
type Req = Serv::Req;
type Resp = Serv::Resp;
async fn serve(
self,
mut ctx: context::Context,
req: Serv::Req,
) -> Result<Serv::Resp, ServerError> {
let ServeThenHook {
serve, mut hook, ..
} = self;
let mut resp = serve.serve(ctx, req).await;
hook.after(&mut ctx, &mut resp).await;
resp
}
}

View File

@@ -0,0 +1,210 @@
// Copyright 2022 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a hook that runs before request execution.
use crate::{context, server::Serve, ServerError};
use futures::prelude::*;
/// A hook that runs before request execution.
#[allow(async_fn_in_trait)]
pub trait BeforeRequest<Req> {
/// The function that is called before request execution.
///
/// If this function returns an error, the request will not be executed and the error will be
/// returned instead.
///
/// This function can also modify the request context. This could be used, for example, to
/// enforce a maximum deadline on all requests.
async fn before(&mut self, ctx: &mut context::Context, req: &Req) -> Result<(), ServerError>;
}
/// A list of hooks that run in order before request execution.
pub trait BeforeRequestList<Req>: BeforeRequest<Req> {
/// The hook returned by `BeforeRequestList::then`.
type Then<Next>: BeforeRequest<Req>
where
Next: BeforeRequest<Req>;
/// Returns a hook that, when run, runs two hooks, first `self` and then `next`.
fn then<Next: BeforeRequest<Req>>(self, next: Next) -> Self::Then<Next>;
/// Same as `then`, but helps the compiler with type inference when Next is a closure.
fn then_fn<
Next: FnMut(&mut context::Context, &Req) -> Fut,
Fut: Future<Output = Result<(), ServerError>>,
>(
self,
next: Next,
) -> Self::Then<Next>
where
Self: Sized,
{
self.then(next)
}
/// The service fn returned by `BeforeRequestList::serving`.
type Serve<S: Serve<Req = Req>>: Serve<Req = Req>;
/// Runs the list of request hooks before execution of the given serve fn.
/// This is equivalent to `serve.before(before_request_chain)` but may be syntactically nicer.
fn serving<S: Serve<Req = Req>>(self, serve: S) -> Self::Serve<S>;
}
impl<F, Fut, Req> BeforeRequest<Req> for F
where
F: FnMut(&mut context::Context, &Req) -> Fut,
Fut: Future<Output = Result<(), ServerError>>,
{
async fn before(&mut self, ctx: &mut context::Context, req: &Req) -> Result<(), ServerError> {
self(ctx, req).await
}
}
/// A Service function that runs a hook before request execution.
#[derive(Clone)]
pub struct HookThenServe<Serv, Hook> {
serve: Serv,
hook: Hook,
}
impl<Serv, Hook> HookThenServe<Serv, Hook> {
pub(crate) fn new(serve: Serv, hook: Hook) -> Self {
Self { serve, hook }
}
}
impl<Serv, Hook> Serve for HookThenServe<Serv, Hook>
where
Serv: Serve,
Hook: BeforeRequest<Serv::Req>,
{
type Req = Serv::Req;
type Resp = Serv::Resp;
async fn serve(
self,
mut ctx: context::Context,
req: Self::Req,
) -> Result<Serv::Resp, ServerError> {
let HookThenServe {
serve, mut hook, ..
} = self;
hook.before(&mut ctx, &req).await?;
serve.serve(ctx, req).await
}
}
/// Returns a request hook builder that runs a series of hooks before request execution.
///
/// Example
///
/// ```rust
/// use futures::{executor::block_on, future};
/// use tarpc::{context, ServerError, server::{Serve, serve, request_hook::{self,
/// BeforeRequest, BeforeRequestList}}};
/// use std::{cell::Cell, io};
///
/// let i = Cell::new(0);
/// let serve = request_hook::before()
/// .then_fn(|_, _| async {
/// assert!(i.get() == 0);
/// i.set(1);
/// Ok(())
/// })
/// .then_fn(|_, _| async {
/// assert!(i.get() == 1);
/// i.set(2);
/// Ok(())
/// })
/// .serving(serve(|_ctx, i| async move { Ok(i + 1) }));
/// let response = serve.clone().serve(context::current(), 1);
/// assert!(block_on(response).is_ok());
/// assert!(i.get() == 2);
/// ```
pub fn before() -> BeforeRequestNil {
BeforeRequestNil
}
/// A list of hooks that run in order before a request is executed.
#[derive(Clone, Copy)]
pub struct BeforeRequestCons<First, Rest>(First, Rest);
/// A noop hook that runs before a request is executed.
#[derive(Clone, Copy)]
pub struct BeforeRequestNil;
impl<Req, First: BeforeRequest<Req>, Rest: BeforeRequest<Req>> BeforeRequest<Req>
for BeforeRequestCons<First, Rest>
{
async fn before(&mut self, ctx: &mut context::Context, req: &Req) -> Result<(), ServerError> {
let BeforeRequestCons(first, rest) = self;
first.before(ctx, req).await?;
rest.before(ctx, req).await?;
Ok(())
}
}
impl<Req> BeforeRequest<Req> for BeforeRequestNil {
async fn before(&mut self, _: &mut context::Context, _: &Req) -> Result<(), ServerError> {
Ok(())
}
}
impl<Req, First: BeforeRequest<Req>, Rest: BeforeRequestList<Req>> BeforeRequestList<Req>
for BeforeRequestCons<First, Rest>
{
type Then<Next> = BeforeRequestCons<First, Rest::Then<Next>> where Next: BeforeRequest<Req>;
fn then<Next: BeforeRequest<Req>>(self, next: Next) -> Self::Then<Next> {
let BeforeRequestCons(first, rest) = self;
BeforeRequestCons(first, rest.then(next))
}
type Serve<S: Serve<Req = Req>> = HookThenServe<S, Self>;
fn serving<S: Serve<Req = Req>>(self, serve: S) -> Self::Serve<S> {
HookThenServe::new(serve, self)
}
}
impl<Req> BeforeRequestList<Req> for BeforeRequestNil {
type Then<Next> = BeforeRequestCons<Next, BeforeRequestNil> where Next: BeforeRequest<Req>;
fn then<Next: BeforeRequest<Req>>(self, next: Next) -> Self::Then<Next> {
BeforeRequestCons(next, BeforeRequestNil)
}
type Serve<S: Serve<Req = Req>> = S;
fn serving<S: Serve<Req = Req>>(self, serve: S) -> S {
serve
}
}
#[test]
fn before_request_list() {
use crate::server::serve;
use futures::executor::block_on;
use std::cell::Cell;
let i = Cell::new(0);
let serve = before()
.then_fn(|_, _| async {
assert!(i.get() == 0);
i.set(1);
Ok(())
})
.then_fn(|_, _| async {
assert!(i.get() == 1);
i.set(2);
Ok(())
})
.serving(serve(|_ctx, i| async move { Ok(i + 1) }));
let response = serve.clone().serve(context::current(), 1);
assert!(block_on(response).is_ok());
assert!(i.get() == 2);
}

View File

@@ -0,0 +1,57 @@
// Copyright 2022 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a hook that runs both before and after request execution.
use super::{after::AfterRequest, before::BeforeRequest};
use crate::{context, server::Serve, ServerError};
use std::marker::PhantomData;
/// A Service function that runs a hook both before and after request execution.
pub struct HookThenServeThenHook<Req, Resp, Serv, Hook> {
serve: Serv,
hook: Hook,
fns: PhantomData<(fn(Req), fn(Resp))>,
}
impl<Req, Resp, Serv, Hook> HookThenServeThenHook<Req, Resp, Serv, Hook> {
pub(crate) fn new(serve: Serv, hook: Hook) -> Self {
Self {
serve,
hook,
fns: PhantomData,
}
}
}
impl<Req, Resp, Serv: Clone, Hook: Clone> Clone for HookThenServeThenHook<Req, Resp, Serv, Hook> {
fn clone(&self) -> Self {
Self {
serve: self.serve.clone(),
hook: self.hook.clone(),
fns: PhantomData,
}
}
}
impl<Req, Resp, Serv, Hook> Serve for HookThenServeThenHook<Req, Resp, Serv, Hook>
where
Serv: Serve<Req = Req, Resp = Resp>,
Hook: BeforeRequest<Req> + AfterRequest<Resp>,
{
type Req = Req;
type Resp = Resp;
async fn serve(self, mut ctx: context::Context, req: Req) -> Result<Serv::Resp, ServerError> {
let HookThenServeThenHook {
serve, mut hook, ..
} = self;
hook.before(&mut ctx, &req).await?;
let mut resp = serve.serve(ctx, req).await;
hook.after(&mut ctx, &mut resp).await;
resp
}
}

View File

@@ -1,16 +1,19 @@
use crate::server::{Channel, Config};
use crate::{context, Request, Response};
use fnv::FnvHashSet;
use futures::{
future::{AbortHandle, AbortRegistration},
task::*,
Sink, Stream,
// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
use crate::{
cancellations::{cancellations, CanceledRequests, RequestCancellation},
context,
server::{Channel, Config, ResponseGuard, TrackedRequest},
Request, Response,
};
use futures::{task::*, Sink, Stream};
use pin_project::pin_project;
use std::collections::VecDeque;
use std::io;
use std::pin::Pin;
use std::time::SystemTime;
use std::{collections::VecDeque, io, pin::Pin, time::SystemTime};
use tracing::Span;
#[pin_project]
pub(crate) struct FakeChannel<In, Out> {
@@ -19,7 +22,9 @@ pub(crate) struct FakeChannel<In, Out> {
#[pin]
pub sink: VecDeque<Out>,
pub config: Config,
pub in_flight_requests: FnvHashSet<u64>,
pub in_flight_requests: super::in_flight_requests::InFlightRequests,
pub request_cancellation: RequestCancellation,
pub canceled_requests: CanceledRequests,
}
impl<In, Out> Stream for FakeChannel<In, Out>
@@ -44,7 +49,7 @@ impl<In, Resp> Sink<Response<Resp>> for FakeChannel<In, Response<Resp>> {
self.as_mut()
.project()
.in_flight_requests
.remove(&response.request_id);
.remove_request(response.request_id);
self.project()
.sink
.start_send(response)
@@ -60,49 +65,61 @@ impl<In, Resp> Sink<Response<Resp>> for FakeChannel<In, Response<Resp>> {
}
}
impl<Req, Resp> Channel for FakeChannel<io::Result<Request<Req>>, Response<Resp>>
impl<Req, Resp> Channel for FakeChannel<io::Result<TrackedRequest<Req>>, Response<Resp>>
where
Req: Unpin,
{
type Req = Req;
type Resp = Resp;
type Transport = ();
fn config(&self) -> &Config {
&self.config
}
fn in_flight_requests(self: Pin<&mut Self>) -> usize {
fn in_flight_requests(&self) -> usize {
self.in_flight_requests.len()
}
fn start_request(self: Pin<&mut Self>, id: u64) -> AbortRegistration {
self.project().in_flight_requests.insert(id);
AbortHandle::new_pair().1
fn transport(&self) -> &() {
&()
}
}
impl<Req, Resp> FakeChannel<io::Result<Request<Req>>, Response<Resp>> {
impl<Req, Resp> FakeChannel<io::Result<TrackedRequest<Req>>, Response<Resp>> {
pub fn push_req(&mut self, id: u64, message: Req) {
self.stream.push_back(Ok(Request {
context: context::Context {
deadline: SystemTime::UNIX_EPOCH,
trace_context: Default::default(),
_non_exhaustive: (),
let (_, abort_registration) = futures::future::AbortHandle::new_pair();
let (request_cancellation, _) = cancellations();
self.stream.push_back(Ok(TrackedRequest {
request: Request {
context: context::Context {
deadline: SystemTime::UNIX_EPOCH,
trace_context: Default::default(),
},
id,
message,
},
abort_registration,
span: Span::none(),
response_guard: ResponseGuard {
request_cancellation,
request_id: id,
cancel: false,
},
id,
message,
_non_exhaustive: (),
}));
}
}
impl FakeChannel<(), ()> {
pub fn default<Req, Resp>() -> FakeChannel<io::Result<Request<Req>>, Response<Resp>> {
pub fn default<Req, Resp>() -> FakeChannel<io::Result<TrackedRequest<Req>>, Response<Resp>> {
let (request_cancellation, canceled_requests) = cancellations();
FakeChannel {
stream: Default::default(),
sink: Default::default(),
config: Default::default(),
in_flight_requests: Default::default(),
request_cancellation,
canceled_requests,
}
}
}
@@ -113,13 +130,10 @@ pub trait PollExt {
impl<T> PollExt for Poll<Option<T>> {
fn is_done(&self) -> bool {
match self {
Poll::Ready(None) => true,
_ => false,
}
matches!(self, Poll::Ready(None))
}
}
pub fn cx() -> Context<'static> {
Context::from_waker(&noop_waker_ref())
Context::from_waker(noop_waker_ref())
}

View File

@@ -16,18 +16,21 @@
//! This crate's design is based on [opencensus
//! tracing](https://opencensus.io/core-concepts/tracing/).
use opentelemetry::trace::TraceContextExt;
use rand::Rng;
use std::{
convert::TryFrom,
fmt::{self, Formatter},
mem,
num::{NonZeroU128, NonZeroU64},
};
use tracing_opentelemetry::OpenTelemetrySpanExt;
/// A context for tracing the execution of processes, distributed or otherwise.
///
/// Consists of a span identifying an event, an optional parent span identifying a causal event
/// that triggered the current span, and a trace with which all related spans are associated.
#[derive(Debug, Default, PartialEq, Eq, Hash, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct Context {
/// An identifier of the trace associated with the current context. A trace ID is typically
/// created at a root span and passed along through all causal events.
@@ -36,33 +39,50 @@ pub struct Context {
/// before making an RPC, and the span ID is sent to the server. The server is free to create
/// its own spans, for which it sets the client's span as the parent span.
pub span_id: SpanId,
/// An identifier of the span that originated the current span. For example, if a server sends
/// an RPC in response to a client request that included a span, the server would create a span
/// for the RPC and set its parent to the span_id in the incoming request's context.
///
/// If `parent_id` is `None`, then this is a root context.
pub parent_id: Option<SpanId>,
/// Indicates whether a sampler has already decided whether or not to sample the trace
/// associated with the Context. If `sampling_decision` is None, then a decision has not yet
/// been made. Downstream samplers do not need to abide by "no sample" decisions--for example,
/// an upstream client may choose to never sample, which may not make sense for the client's
/// dependencies. On the other hand, if an upstream process has chosen to sample this trace,
/// then the downstream samplers are expected to respect that decision and also sample the
/// trace. Otherwise, the full trace would not be able to be reconstructed.
pub sampling_decision: SamplingDecision,
}
/// A 128-bit UUID identifying a trace. All spans caused by the same originating span share the
/// same trace ID.
#[derive(Debug, Default, PartialEq, Eq, Hash, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct TraceId(u128);
#[derive(Default, PartialEq, Eq, Hash, Clone, Copy)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct TraceId(#[cfg_attr(feature = "serde1", serde(with = "u128_serde"))] u128);
/// A 64-bit identifier of a span within a trace. The identifier is unique within the span's trace.
#[derive(Debug, Default, PartialEq, Eq, Hash, Clone, Copy)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Default, PartialEq, Eq, Hash, Clone, Copy)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct SpanId(u64);
/// Indicates whether a sampler has decided whether or not to sample the trace associated with the
/// Context. Downstream samplers do not need to abide by "no sample" decisions--for example, an
/// upstream client may choose to never sample, which may not make sense for the client's
/// dependencies. On the other hand, if an upstream process has chosen to sample this trace, then
/// the downstream samplers are expected to respect that decision and also sample the trace.
/// Otherwise, the full trace would not be able to be reconstructed reliably.
#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
#[repr(u8)]
pub enum SamplingDecision {
/// The associated span was sampled by its creating process. Child spans must also be sampled.
Sampled,
/// The associated span was not sampled by its creating process.
Unsampled,
}
impl Context {
/// Constructs a new root context. A root context is one with no parent span.
pub fn new_root() -> Self {
let rng = &mut rand::thread_rng();
Context {
trace_id: TraceId::random(rng),
span_id: SpanId::random(rng),
parent_id: None,
/// Constructs a new context with the trace ID and sampling decision inherited from the parent.
pub(crate) fn new_child(&self) -> Self {
Self {
trace_id: self.trace_id,
span_id: SpanId::random(&mut rand::thread_rng()),
sampling_decision: self.sampling_decision,
}
}
}
@@ -71,17 +91,128 @@ impl TraceId {
/// Returns a random trace ID that can be assumed to be globally unique if `rng` generates
/// actually-random numbers.
pub fn random<R: Rng>(rng: &mut R) -> Self {
TraceId(u128::from(rng.next_u64()) << mem::size_of::<u64>() | u128::from(rng.next_u64()))
TraceId(rng.gen::<NonZeroU128>().get())
}
/// Returns true iff the trace ID is 0.
pub fn is_none(&self) -> bool {
self.0 == 0
}
}
impl SpanId {
/// Returns a random span ID that can be assumed to be unique within a single trace.
pub fn random<R: Rng>(rng: &mut R) -> Self {
SpanId(rng.next_u64())
SpanId(rng.gen::<NonZeroU64>().get())
}
/// Returns true iff the span ID is 0.
pub fn is_none(&self) -> bool {
self.0 == 0
}
}
impl From<TraceId> for u128 {
fn from(trace_id: TraceId) -> Self {
trace_id.0
}
}
impl From<u128> for TraceId {
fn from(trace_id: u128) -> Self {
Self(trace_id)
}
}
impl From<SpanId> for u64 {
fn from(span_id: SpanId) -> Self {
span_id.0
}
}
impl From<u64> for SpanId {
fn from(span_id: u64) -> Self {
Self(span_id)
}
}
impl From<opentelemetry::trace::TraceId> for TraceId {
fn from(trace_id: opentelemetry::trace::TraceId) -> Self {
Self::from(u128::from_be_bytes(trace_id.to_bytes()))
}
}
impl From<TraceId> for opentelemetry::trace::TraceId {
fn from(trace_id: TraceId) -> Self {
Self::from_bytes(u128::from(trace_id).to_be_bytes())
}
}
impl From<opentelemetry::trace::SpanId> for SpanId {
fn from(span_id: opentelemetry::trace::SpanId) -> Self {
Self::from(u64::from_be_bytes(span_id.to_bytes()))
}
}
impl From<SpanId> for opentelemetry::trace::SpanId {
fn from(span_id: SpanId) -> Self {
Self::from_bytes(u64::from(span_id).to_be_bytes())
}
}
impl TryFrom<&tracing::Span> for Context {
type Error = NoActiveSpan;
fn try_from(span: &tracing::Span) -> Result<Self, NoActiveSpan> {
let context = span.context();
if context.has_active_span() {
Ok(Self::from(context.span()))
} else {
Err(NoActiveSpan)
}
}
}
impl From<opentelemetry::trace::SpanRef<'_>> for Context {
fn from(span: opentelemetry::trace::SpanRef<'_>) -> Self {
let otel_ctx = span.span_context();
Self {
trace_id: TraceId::from(otel_ctx.trace_id()),
span_id: SpanId::from(otel_ctx.span_id()),
sampling_decision: SamplingDecision::from(otel_ctx),
}
}
}
impl From<SamplingDecision> for opentelemetry::trace::TraceFlags {
fn from(decision: SamplingDecision) -> Self {
match decision {
SamplingDecision::Sampled => opentelemetry::trace::TraceFlags::SAMPLED,
SamplingDecision::Unsampled => opentelemetry::trace::TraceFlags::default(),
}
}
}
impl From<&opentelemetry::trace::SpanContext> for SamplingDecision {
fn from(context: &opentelemetry::trace::SpanContext) -> Self {
if context.is_sampled() {
SamplingDecision::Sampled
} else {
SamplingDecision::Unsampled
}
}
}
impl Default for SamplingDecision {
fn default() -> Self {
Self::Unsampled
}
}
/// Returned when a [`Context`] cannot be constructed from a [`Span`](tracing::Span).
#[derive(Debug)]
pub struct NoActiveSpan;
impl fmt::Display for TraceId {
fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
write!(f, "{:02x}", self.0)?;
@@ -89,9 +220,42 @@ impl fmt::Display for TraceId {
}
}
impl fmt::Debug for TraceId {
fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
write!(f, "{:02x}", self.0)?;
Ok(())
}
}
impl fmt::Display for SpanId {
fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
write!(f, "{:02x}", self.0)?;
Ok(())
}
}
impl fmt::Debug for SpanId {
fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
write!(f, "{:02x}", self.0)?;
Ok(())
}
}
#[cfg(feature = "serde1")]
mod u128_serde {
pub fn serialize<S>(u: &u128, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serde::Serialize::serialize(&u.to_le_bytes(), serializer)
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<u128, D::Error>
where
D: serde::Deserializer<'de>,
{
Ok(u128::from_le_bytes(serde::Deserialize::deserialize(
deserializer,
)?))
}
}

40
tarpc/src/transport.rs Normal file
View File

@@ -0,0 +1,40 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Provides a [`Transport`](sealed::Transport) trait as well as implementations.
//!
//! The rpc crate is transport- and protocol-agnostic. Any transport that impls [`Transport`](sealed::Transport)
//! can be plugged in, using whatever protocol it wants.
pub mod channel;
pub(crate) mod sealed {
use futures::prelude::*;
use std::error::Error;
/// A bidirectional stream ([`Sink`] + [`Stream`]) of messages.
pub trait Transport<SinkItem, Item>
where
Self: Stream<Item = Result<Item, <Self as Sink<SinkItem>>::Error>>,
Self: Sink<SinkItem, Error = <Self as Transport<SinkItem, Item>>::TransportError>,
<Self as Sink<SinkItem>>::Error: Error,
{
/// Associated type where clauses are not elaborated; this associated type allows users
/// bounding types by Transport to avoid having to explicitly add `T::Error: Error` to their
/// bounds.
type TransportError: Error + Send + Sync + 'static;
}
impl<T, SinkItem, Item, E> Transport<SinkItem, Item> for T
where
T: ?Sized,
T: Stream<Item = Result<Item, E>>,
T: Sink<SinkItem, Error = E>,
T::Error: Error + Send + Sync + 'static,
{
type TransportError = E;
}
}

View File

@@ -0,0 +1,219 @@
// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.
//! Transports backed by in-memory channels.
use futures::{task::*, Sink, Stream};
use pin_project::pin_project;
use std::{error::Error, pin::Pin};
use tokio::sync::mpsc;
/// Errors that occur in the sending or receiving of messages over a channel.
#[derive(thiserror::Error, Debug)]
pub enum ChannelError {
/// An error occurred readying to send into the channel.
#[error("an error occurred readying to send into the channel")]
Ready(#[source] Box<dyn Error + Send + Sync + 'static>),
/// An error occurred sending into the channel.
#[error("an error occurred sending into the channel")]
Send(#[source] Box<dyn Error + Send + Sync + 'static>),
/// An error occurred receiving from the channel.
#[error("an error occurred receiving from the channel")]
Receive(#[source] Box<dyn Error + Send + Sync + 'static>),
}
/// Returns two unbounded channel peers. Each [`Stream`] yields items sent through the other's
/// [`Sink`].
pub fn unbounded<SinkItem, Item>() -> (
UnboundedChannel<SinkItem, Item>,
UnboundedChannel<Item, SinkItem>,
) {
let (tx1, rx2) = mpsc::unbounded_channel();
let (tx2, rx1) = mpsc::unbounded_channel();
(
UnboundedChannel { tx: tx1, rx: rx1 },
UnboundedChannel { tx: tx2, rx: rx2 },
)
}
/// A bi-directional channel backed by an [`UnboundedSender`](mpsc::UnboundedSender)
/// and [`UnboundedReceiver`](mpsc::UnboundedReceiver).
#[derive(Debug)]
pub struct UnboundedChannel<Item, SinkItem> {
rx: mpsc::UnboundedReceiver<Item>,
tx: mpsc::UnboundedSender<SinkItem>,
}
impl<Item, SinkItem> Stream for UnboundedChannel<Item, SinkItem> {
type Item = Result<Item, ChannelError>;
fn poll_next(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Item, ChannelError>>> {
self.rx
.poll_recv(cx)
.map(|option| option.map(Ok))
.map_err(ChannelError::Receive)
}
}
const CLOSED_MESSAGE: &str = "the channel is closed and cannot accept new items for sending";
impl<Item, SinkItem> Sink<SinkItem> for UnboundedChannel<Item, SinkItem> {
type Error = ChannelError;
fn poll_ready(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(if self.tx.is_closed() {
Err(ChannelError::Ready(CLOSED_MESSAGE.into()))
} else {
Ok(())
})
}
fn start_send(self: Pin<&mut Self>, item: SinkItem) -> Result<(), Self::Error> {
self.tx
.send(item)
.map_err(|_| ChannelError::Send(CLOSED_MESSAGE.into()))
}
fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
// UnboundedSender requires no flushing.
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
// UnboundedSender can't initiate closure.
Poll::Ready(Ok(()))
}
}
/// Returns two channel peers with buffer equal to `capacity`. Each [`Stream`] yields items sent
/// through the other's [`Sink`].
pub fn bounded<SinkItem, Item>(
capacity: usize,
) -> (Channel<SinkItem, Item>, Channel<Item, SinkItem>) {
let (tx1, rx2) = futures::channel::mpsc::channel(capacity);
let (tx2, rx1) = futures::channel::mpsc::channel(capacity);
(Channel { tx: tx1, rx: rx1 }, Channel { tx: tx2, rx: rx2 })
}
/// A bi-directional channel backed by a [`Sender`](futures::channel::mpsc::Sender)
/// and [`Receiver`](futures::channel::mpsc::Receiver).
#[pin_project]
#[derive(Debug)]
pub struct Channel<Item, SinkItem> {
#[pin]
rx: futures::channel::mpsc::Receiver<Item>,
#[pin]
tx: futures::channel::mpsc::Sender<SinkItem>,
}
impl<Item, SinkItem> Stream for Channel<Item, SinkItem> {
type Item = Result<Item, ChannelError>;
fn poll_next(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Item, ChannelError>>> {
self.project()
.rx
.poll_next(cx)
.map(|option| option.map(Ok))
.map_err(ChannelError::Receive)
}
}
impl<Item, SinkItem> Sink<SinkItem> for Channel<Item, SinkItem> {
type Error = ChannelError;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project()
.tx
.poll_ready(cx)
.map_err(|e| ChannelError::Ready(Box::new(e)))
}
fn start_send(self: Pin<&mut Self>, item: SinkItem) -> Result<(), Self::Error> {
self.project()
.tx
.start_send(item)
.map_err(|e| ChannelError::Send(Box::new(e)))
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project()
.tx
.poll_flush(cx)
.map_err(|e| ChannelError::Send(Box::new(e)))
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project()
.tx
.poll_close(cx)
.map_err(|e| ChannelError::Send(Box::new(e)))
}
}
#[cfg(all(test, feature = "tokio1"))]
mod tests {
use crate::{
client::{self, RpcError},
context,
server::{incoming::Incoming, serve, BaseChannel},
transport::{
self,
channel::{Channel, UnboundedChannel},
},
ServerError,
};
use assert_matches::assert_matches;
use futures::{prelude::*, stream};
use std::io;
use tracing::trace;
#[test]
fn ensure_is_transport() {
fn is_transport<SinkItem, Item, T: crate::Transport<SinkItem, Item>>() {}
is_transport::<(), (), UnboundedChannel<(), ()>>();
is_transport::<(), (), Channel<(), ()>>();
}
#[tokio::test]
async fn integration() -> anyhow::Result<()> {
let _ = tracing_subscriber::fmt::try_init();
let (client_channel, server_channel) = transport::channel::unbounded();
tokio::spawn(
stream::once(future::ready(server_channel))
.map(BaseChannel::with_defaults)
.execute(serve(|_ctx, request: String| async move {
request.parse::<u64>().map_err(|_| {
ServerError::new(
io::ErrorKind::InvalidInput,
format!("{request:?} is not an int"),
)
})
}))
.for_each(|channel| async move {
tokio::spawn(channel.for_each(|response| response));
}),
);
let client = client::new(client::Config::default(), client_channel).spawn();
let response1 = client.call(context::current(), "", "123".into()).await;
let response2 = client.call(context::current(), "", "abc".into()).await;
trace!("response1: {:?}, response2: {:?}", response1, response2);
assert_matches!(response1, Ok(123));
assert_matches!(response2, Err(RpcError::Server(e)) if e.kind == io::ErrorKind::InvalidInput);
Ok(())
}
}

View File

@@ -10,7 +10,8 @@ use std::{
time::{Duration, SystemTime},
};
#[cfg(feature = "serde")]
#[cfg(feature = "serde1")]
#[cfg_attr(docsrs, doc(cfg(feature = "serde1")))]
pub mod serde;
/// Extension trait for [SystemTimes](SystemTime) in the future, i.e. deadlines.
@@ -37,11 +38,34 @@ where
H: BuildHasher,
{
fn compact(&mut self, usage_ratio_threshold: f64) {
if self.capacity() > 1000 {
let usage_ratio = self.len() as f64 / self.capacity() as f64;
if usage_ratio < usage_ratio_threshold {
self.shrink_to_fit();
}
}
let usage_ratio_threshold = usage_ratio_threshold.clamp(f64::MIN_POSITIVE, 1.);
let cap = f64::max(1000., self.len() as f64 / usage_ratio_threshold);
self.shrink_to(cap as usize);
}
}
#[test]
fn test_compact() {
let mut map = HashMap::with_capacity(2048);
assert_eq!(map.capacity(), 3584);
// Make usage ratio 25%
for i in 0..896 {
map.insert(format!("k{i}"), "v");
}
map.compact(-1.0);
assert_eq!(map.capacity(), 3584);
map.compact(0.25);
assert_eq!(map.capacity(), 3584);
map.compact(0.50);
assert_eq!(map.capacity(), 1792);
map.compact(1.0);
assert_eq!(map.capacity(), 1792);
map.compact(2.0);
assert_eq!(map.capacity(), 1792);
}

View File

@@ -5,30 +5,7 @@
// https://opensource.org/licenses/MIT.
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::{
io,
time::{Duration, SystemTime},
};
/// Serializes `system_time` as a `u64` equal to the number of seconds since the epoch.
pub fn serialize_epoch_secs<S>(system_time: &SystemTime, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
system_time
.duration_since(SystemTime::UNIX_EPOCH)
.unwrap_or(Duration::from_secs(0))
.as_secs() // Only care about second precision
.serialize(serializer)
}
/// Deserializes [`SystemTime`] from a `u64` equal to the number of seconds since the epoch.
pub fn deserialize_epoch_secs<'de, D>(deserializer: D) -> Result<SystemTime, D::Error>
where
D: Deserializer<'de>,
{
Ok(SystemTime::UNIX_EPOCH + Duration::from_secs(u64::deserialize(deserializer)?))
}
use std::io;
/// Serializes [`io::ErrorKind`] as a `u32`.
#[allow(clippy::trivially_copy_pass_by_ref)] // Exact fn signature required by serde derive

View File

@@ -0,0 +1,7 @@
#[test]
fn ui() {
let t = trybuild::TestCases::new();
t.compile_fail("tests/compile_fail/*.rs");
#[cfg(all(feature = "serde-transport", feature = "tcp"))]
t.compile_fail("tests/compile_fail/serde_transport/*.rs");
}

View File

@@ -0,0 +1,15 @@
use tarpc::client;
#[tarpc::service]
trait World {
async fn hello(name: String) -> String;
}
fn main() {
let (client_transport, _) = tarpc::transport::channel::unbounded();
#[deny(unused_must_use)]
{
WorldClient::new(client::Config::default(), client_transport).dispatch;
}
}

View File

@@ -0,0 +1,15 @@
error: unused `RequestDispatch` that must be used
--> tests/compile_fail/must_use_request_dispatch.rs:13:9
|
13 | WorldClient::new(client::Config::default(), client_transport).dispatch;
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
note: the lint level is defined here
--> tests/compile_fail/must_use_request_dispatch.rs:11:12
|
11 | #[deny(unused_must_use)]
| ^^^^^^^^^^^^^^^
help: use `let _ = ...` to ignore the resulting value
|
13 | let _ = WorldClient::new(client::Config::default(), client_transport).dispatch;
| +++++++

View File

@@ -0,0 +1,9 @@
use tarpc::serde_transport;
use tokio_serde::formats::Json;
fn main() {
#[deny(unused_must_use)]
{
serde_transport::tcp::connect::<_, (), (), _, _>("0.0.0.0:0", Json::default);
}
}

View File

@@ -0,0 +1,15 @@
error: unused `tarpc::serde_transport::tcp::Connect` that must be used
--> tests/compile_fail/serde_transport/must_use_tcp_connect.rs:7:9
|
7 | serde_transport::tcp::connect::<_, (), (), _, _>("0.0.0.0:0", Json::default);
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
note: the lint level is defined here
--> tests/compile_fail/serde_transport/must_use_tcp_connect.rs:5:12
|
5 | #[deny(unused_must_use)]
| ^^^^^^^^^^^^^^^
help: use `let _ = ...` to ignore the resulting value
|
7 | let _ = serde_transport::tcp::connect::<_, (), (), _, _>("0.0.0.0:0", Json::default);
| +++++++

View File

@@ -0,0 +1,6 @@
#[tarpc::service]
trait World {
async fn pat((a, b): (u8, u32));
}
fn main() {}

View File

@@ -0,0 +1,5 @@
error: patterns aren't allowed in RPC args
--> $DIR/tarpc_service_arg_pat.rs:3:18
|
3 | async fn pat((a, b): (u8, u32));
| ^^^^^^

View File

@@ -0,0 +1,6 @@
#[tarpc::service]
trait World {
async fn new();
}
fn main() {}

View File

@@ -0,0 +1,5 @@
error: method name conflicts with generated fn `WorldClient::new`
--> $DIR/tarpc_service_fn_new.rs:3:14
|
3 | async fn new();
| ^^^

View File

@@ -0,0 +1,6 @@
#[tarpc::service]
trait World {
async fn serve();
}
fn main() {}

View File

@@ -0,0 +1,5 @@
error: method name conflicts with generated fn `World::serve`
--> $DIR/tarpc_service_fn_serve.rs:3:14
|
3 | async fn serve();
| ^^^^^

View File

@@ -0,0 +1,61 @@
use futures::prelude::*;
use tarpc::serde_transport;
use tarpc::{
client, context,
server::{incoming::Incoming, BaseChannel},
};
use tokio_serde::formats::Json;
#[tarpc::derive_serde]
#[derive(Debug, PartialEq, Eq)]
pub enum TestData {
Black,
White,
}
#[tarpc::service]
pub trait ColorProtocol {
async fn get_opposite_color(color: TestData) -> TestData;
}
#[derive(Clone)]
struct ColorServer;
impl ColorProtocol for ColorServer {
async fn get_opposite_color(self, _: context::Context, color: TestData) -> TestData {
match color {
TestData::White => TestData::Black,
TestData::Black => TestData::White,
}
}
}
#[cfg(test)]
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
#[tokio::test]
async fn test_call() -> anyhow::Result<()> {
let transport = tarpc::serde_transport::tcp::listen("localhost:56797", Json::default).await?;
let addr = transport.local_addr();
tokio::spawn(
transport
.take(1)
.filter_map(|r| async { r.ok() })
.map(BaseChannel::with_defaults)
.execute(ColorServer.serve())
.map(|channel| channel.for_each(spawn))
.for_each(spawn),
);
let transport = serde_transport::tcp::connect(addr, Json::default).await?;
let client = ColorProtocolClient::new(client::Config::default(), transport).spawn();
let color = client
.get_opposite_color(context::current(), TestData::White)
.await?;
assert_eq!(color, TestData::Black);
Ok(())
}

View File

@@ -0,0 +1,17 @@
#![no_implicit_prelude]
extern crate tarpc as some_random_other_name;
#[cfg(feature = "serde1")]
mod serde1_feature {
#[::tarpc::derive_serde]
#[derive(Debug, PartialEq, Eq)]
pub enum TestData {
Black,
White,
}
}
#[::tarpc::service]
pub trait ColorProtocol {
async fn get_opposite_color(color: u8) -> u8;
}

View File

@@ -1,16 +1,16 @@
use assert_matches::assert_matches;
use futures::{
future::{ready, Ready},
future::{join_all, ready},
prelude::*,
};
use std::io;
use std::time::{Duration, SystemTime};
use tarpc::{
client::{self},
context, serde_transport,
server::{self, BaseChannel, Channel, Handler},
context,
server::{incoming::Incoming, BaseChannel, Channel},
transport::channel,
};
use tokio_serde::formats::Json;
use tokio::join;
#[tarpc_plugins::service]
trait Service {
@@ -22,56 +22,101 @@ trait Service {
struct Server;
impl Service for Server {
type AddFut = Ready<i32>;
fn add(self, _: context::Context, x: i32, y: i32) -> Self::AddFut {
ready(x + y)
async fn add(self, _: context::Context, x: i32, y: i32) -> i32 {
x + y
}
type HeyFut = Ready<String>;
fn hey(self, _: context::Context, name: String) -> Self::HeyFut {
ready(format!("Hey, {}.", name))
async fn hey(self, _: context::Context, name: String) -> String {
format!("Hey, {name}.")
}
}
#[tokio::test(threaded_scheduler)]
async fn sequential() -> io::Result<()> {
let _ = env_logger::try_init();
#[tokio::test]
async fn sequential() {
let (tx, rx) = tarpc::transport::channel::unbounded();
let client = client::new(client::Config::default(), tx).spawn();
let channel = BaseChannel::with_defaults(rx);
tokio::spawn(
channel
.execute(tarpc::server::serve(|_, i| async move { Ok(i + 1) }))
.for_each(|response| response),
);
assert_eq!(
client.call(context::current(), "AddOne", 1).await.unwrap(),
2
);
}
#[tokio::test]
async fn dropped_channel_aborts_in_flight_requests() -> anyhow::Result<()> {
#[tarpc_plugins::service]
trait Loop {
async fn r#loop();
}
#[derive(Clone)]
struct LoopServer;
#[derive(Debug)]
struct AllHandlersComplete;
impl Loop for LoopServer {
async fn r#loop(self, _: context::Context) {
loop {
futures::pending!();
}
}
}
let _ = tracing_subscriber::fmt::try_init();
let (tx, rx) = channel::unbounded();
tokio::spawn(
BaseChannel::new(server::Config::default(), rx)
.respond_with(Server.serve())
.execute(),
);
// Set up a client that initiates a long-lived request.
// The request will complete in error when the server drops the connection.
tokio::spawn(async move {
let client = LoopClient::new(client::Config::default(), tx).spawn();
let mut client = ServiceClient::new(client::Config::default(), tx).spawn()?;
let mut ctx = context::current();
ctx.deadline = SystemTime::now() + Duration::from_secs(60 * 60);
let _ = client.r#loop(ctx).await;
});
assert_matches!(client.add(context::current(), 1, 2).await, Ok(3));
assert_matches!(
client.hey(context::current(), "Tim".into()).await,
Ok(ref s) if s == "Hey, Tim.");
let mut requests = BaseChannel::with_defaults(rx).requests();
// Reading a request should trigger the request being registered with BaseChannel.
let first_request = requests.next().await.unwrap()?;
// Dropping the channel should trigger cleanup of outstanding requests.
drop(requests);
// In-flight requests should be aborted by channel cleanup.
// The first and only request sent by the client is `loop`, which is an infinite loop
// on the server side, so if cleanup was not triggered, this line should hang indefinitely.
first_request.execute(LoopServer.serve()).await;
Ok(())
}
#[cfg(feature = "serde1")]
#[tokio::test(threaded_scheduler)]
async fn serde() -> io::Result<()> {
let _ = env_logger::try_init();
#[cfg(all(feature = "serde-transport", feature = "tcp"))]
#[tokio::test]
async fn serde_tcp() -> anyhow::Result<()> {
use tarpc::serde_transport;
use tokio_serde::formats::Json;
let transport = serde_transport::tcp::listen("localhost:56789", Json::default).await?;
let _ = tracing_subscriber::fmt::try_init();
let transport = tarpc::serde_transport::tcp::listen("localhost:56789", Json::default).await?;
let addr = transport.local_addr();
tokio::spawn(
tarpc::Server::default()
.incoming(transport.take(1).filter_map(|r| async { r.ok() }))
.respond_with(Server.serve()),
transport
.take(1)
.filter_map(|r| async { r.ok() })
.map(BaseChannel::with_defaults)
.execute(Server.serve())
.map(|channel| channel.for_each(spawn))
.for_each(spawn),
);
let transport = serde_transport::tcp::connect(addr, Json::default()).await?;
let mut client = ServiceClient::new(client::Config::default(), transport).spawn()?;
let transport = serde_transport::tcp::connect(addr, Json::default).await?;
let client = ServiceClient::new(client::Config::default(), transport).spawn();
assert_matches!(client.add(context::current(), 1, 2).await, Ok(3));
assert_matches!(
@@ -82,27 +127,57 @@ async fn serde() -> io::Result<()> {
Ok(())
}
#[tokio::test(threaded_scheduler)]
async fn concurrent() -> io::Result<()> {
let _ = env_logger::try_init();
#[cfg(all(feature = "serde-transport", feature = "unix", unix))]
#[tokio::test]
async fn serde_uds() -> anyhow::Result<()> {
use tarpc::serde_transport;
use tokio_serde::formats::Json;
let _ = tracing_subscriber::fmt::try_init();
let sock = tarpc::serde_transport::unix::TempPathBuf::with_random("uds");
let transport = tarpc::serde_transport::unix::listen(&sock, Json::default).await?;
tokio::spawn(
transport
.take(1)
.filter_map(|r| async { r.ok() })
.map(BaseChannel::with_defaults)
.execute(Server.serve())
.map(|channel| channel.for_each(spawn))
.for_each(spawn),
);
let transport = serde_transport::unix::connect(&sock, Json::default).await?;
let client = ServiceClient::new(client::Config::default(), transport).spawn();
// Save results using socket so we can clean the socket even if our test assertions fail
let res1 = client.add(context::current(), 1, 2).await;
let res2 = client.hey(context::current(), "Tim".to_string()).await;
assert_matches!(res1, Ok(3));
assert_matches!(res2, Ok(ref s) if s == "Hey, Tim.");
Ok(())
}
#[tokio::test]
async fn concurrent() -> anyhow::Result<()> {
let _ = tracing_subscriber::fmt::try_init();
let (tx, rx) = channel::unbounded();
tokio::spawn(
tarpc::Server::default()
.incoming(stream::once(ready(rx)))
.respond_with(Server.serve()),
stream::once(ready(rx))
.map(BaseChannel::with_defaults)
.execute(Server.serve())
.map(|channel| channel.for_each(spawn))
.for_each(spawn),
);
let client = ServiceClient::new(client::Config::default(), tx).spawn()?;
let client = ServiceClient::new(client::Config::default(), tx).spawn();
let mut c = client.clone();
let req1 = c.add(context::current(), 1, 2);
let mut c = client.clone();
let req2 = c.add(context::current(), 3, 4);
let mut c = client.clone();
let req3 = c.hey(context::current(), "Tim".to_string());
let req1 = client.add(context::current(), 1, 2);
let req2 = client.add(context::current(), 3, 4);
let req3 = client.hey(context::current(), "Tim".to_string());
assert_matches!(req1.await, Ok(3));
assert_matches!(req2.await, Ok(7));
@@ -110,3 +185,91 @@ async fn concurrent() -> io::Result<()> {
Ok(())
}
#[tokio::test]
async fn concurrent_join() -> anyhow::Result<()> {
let _ = tracing_subscriber::fmt::try_init();
let (tx, rx) = channel::unbounded();
tokio::spawn(
stream::once(ready(rx))
.map(BaseChannel::with_defaults)
.execute(Server.serve())
.map(|channel| channel.for_each(spawn))
.for_each(spawn),
);
let client = ServiceClient::new(client::Config::default(), tx).spawn();
let req1 = client.add(context::current(), 1, 2);
let req2 = client.add(context::current(), 3, 4);
let req3 = client.hey(context::current(), "Tim".to_string());
let (resp1, resp2, resp3) = join!(req1, req2, req3);
assert_matches!(resp1, Ok(3));
assert_matches!(resp2, Ok(7));
assert_matches!(resp3, Ok(ref s) if s == "Hey, Tim.");
Ok(())
}
#[cfg(test)]
async fn spawn(fut: impl Future<Output = ()> + Send + 'static) {
tokio::spawn(fut);
}
#[tokio::test]
async fn concurrent_join_all() -> anyhow::Result<()> {
let _ = tracing_subscriber::fmt::try_init();
let (tx, rx) = channel::unbounded();
tokio::spawn(
BaseChannel::with_defaults(rx)
.execute(Server.serve())
.for_each(spawn),
);
let client = ServiceClient::new(client::Config::default(), tx).spawn();
let req1 = client.add(context::current(), 1, 2);
let req2 = client.add(context::current(), 3, 4);
let responses = join_all(vec![req1, req2]).await;
assert_matches!(responses[0], Ok(3));
assert_matches!(responses[1], Ok(7));
Ok(())
}
#[tokio::test]
async fn counter() -> anyhow::Result<()> {
#[tarpc::service]
trait Counter {
async fn count() -> u32;
}
struct CountService(u32);
impl Counter for &mut CountService {
async fn count(self, _: context::Context) -> u32 {
self.0 += 1;
self.0
}
}
let (tx, rx) = channel::unbounded();
tokio::spawn(async {
let mut requests = BaseChannel::with_defaults(rx).requests();
let mut counter = CountService(0);
while let Some(Ok(request)) = requests.next().await {
request.execute(counter.serve()).await;
}
});
let client = CounterClient::new(client::Config::default(), tx).spawn();
assert_matches!(client.count(context::current()).await, Ok(1));
assert_matches!(client.count(context::current()).await, Ok(2));
Ok(())
}