AsyncUDPServerSocket passes socket in callback
[folly.git] / folly / wangle / bootstrap / ServerBootstrap-inl.h
1 /*
2  * Copyright 2015 Facebook, Inc.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *   http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 #pragma once
17
18 #include <folly/wangle/acceptor/Acceptor.h>
19 #include <folly/wangle/bootstrap/ServerSocketFactory.h>
20 #include <folly/io/async/EventBaseManager.h>
21 #include <folly/wangle/concurrent/IOThreadPoolExecutor.h>
22 #include <folly/wangle/acceptor/ManagedConnection.h>
23 #include <folly/wangle/channel/ChannelPipeline.h>
24 #include <folly/wangle/channel/ChannelHandler.h>
25
26 namespace folly {
27
28 template <typename Pipeline>
29 class ServerAcceptor
30     : public Acceptor
31     , public folly::wangle::ChannelHandlerAdapter<void*, std::exception> {
32   typedef std::unique_ptr<Pipeline,
33                           folly::DelayedDestruction::Destructor> PipelinePtr;
34
35   class ServerConnection : public wangle::ManagedConnection {
36    public:
37     explicit ServerConnection(PipelinePtr pipeline)
38         : pipeline_(std::move(pipeline)) {}
39
40     ~ServerConnection() {
41     }
42
43     void timeoutExpired() noexcept {
44     }
45
46     void describe(std::ostream& os) const {}
47     bool isBusy() const {
48       return false;
49     }
50     void notifyPendingShutdown() {}
51     void closeWhenIdle() {}
52     void dropConnection() {
53       delete this;
54     }
55     void dumpConnectionState(uint8_t loglevel) {}
56    private:
57     PipelinePtr pipeline_;
58   };
59
60  public:
61   explicit ServerAcceptor(
62         std::shared_ptr<PipelineFactory<Pipeline>> pipelineFactory,
63         std::shared_ptr<folly::wangle::ChannelPipeline<
64                           void*, std::exception>> acceptorPipeline,
65         EventBase* base)
66       : Acceptor(ServerSocketConfig())
67       , base_(base)
68       , childPipelineFactory_(pipelineFactory)
69       , acceptorPipeline_(acceptorPipeline) {
70     Acceptor::init(nullptr, base_);
71     CHECK(acceptorPipeline_);
72
73     acceptorPipeline_->addBack(folly::wangle::ChannelHandlerPtr<ServerAcceptor, false>(this));
74     acceptorPipeline_->finalize();
75   }
76
77   void read(Context* ctx, void* conn) {
78     AsyncSocket::UniquePtr transport((AsyncSocket*)conn);
79       std::unique_ptr<Pipeline,
80                        folly::DelayedDestruction::Destructor>
81       pipeline(childPipelineFactory_->newPipeline(
82         std::shared_ptr<AsyncSocket>(
83           transport.release(),
84           folly::DelayedDestruction::Destructor())));
85     auto connection = new ServerConnection(std::move(pipeline));
86     Acceptor::addConnection(connection);
87   }
88
89   folly::Future<void> write(Context* ctx, std::exception e) {
90     return ctx->fireWrite(e);
91   }
92
93   /* See Acceptor::onNewConnection for details */
94   void onNewConnection(
95     AsyncSocket::UniquePtr transport, const SocketAddress* address,
96     const std::string& nextProtocolName, const TransportInfo& tinfo) {
97     acceptorPipeline_->read(transport.release());
98   }
99
100   // UDP thunk
101   void onDataAvailable(std::shared_ptr<AsyncUDPSocket> socket,
102                        const folly::SocketAddress& addr,
103                        std::unique_ptr<folly::IOBuf> buf,
104                        bool truncated) noexcept {
105     acceptorPipeline_->read(buf.release());
106   }
107
108  private:
109   EventBase* base_;
110
111   std::shared_ptr<PipelineFactory<Pipeline>> childPipelineFactory_;
112   std::shared_ptr<folly::wangle::ChannelPipeline<
113     void*, std::exception>> acceptorPipeline_;
114 };
115
116 template <typename Pipeline>
117 class ServerAcceptorFactory : public AcceptorFactory {
118  public:
119   explicit ServerAcceptorFactory(
120     std::shared_ptr<PipelineFactory<Pipeline>> factory,
121     std::shared_ptr<PipelineFactory<folly::wangle::ChannelPipeline<
122     void*, std::exception>>> pipeline)
123     : factory_(factory)
124     , pipeline_(pipeline) {}
125
126   std::shared_ptr<Acceptor> newAcceptor(EventBase* base) {
127     std::shared_ptr<folly::wangle::ChannelPipeline<
128                       void*, std::exception>> pipeline(
129                         pipeline_->newPipeline(nullptr));
130     return std::make_shared<ServerAcceptor<Pipeline>>(factory_, pipeline, base);
131   }
132  private:
133   std::shared_ptr<PipelineFactory<Pipeline>> factory_;
134   std::shared_ptr<PipelineFactory<
135     folly::wangle::ChannelPipeline<
136       void*, std::exception>>> pipeline_;
137 };
138
139 class ServerWorkerPool : public folly::wangle::ThreadPoolExecutor::Observer {
140  public:
141   explicit ServerWorkerPool(
142     std::shared_ptr<AcceptorFactory> acceptorFactory,
143     folly::wangle::IOThreadPoolExecutor* exec,
144     std::shared_ptr<std::vector<std::shared_ptr<folly::AsyncSocketBase>>> sockets,
145     std::shared_ptr<ServerSocketFactory> socketFactory)
146       : acceptorFactory_(acceptorFactory)
147       , exec_(exec)
148       , sockets_(sockets)
149       , socketFactory_(socketFactory) {
150     CHECK(exec);
151   }
152
153   template <typename F>
154   void forEachWorker(F&& f) const;
155
156   void threadStarted(
157     folly::wangle::ThreadPoolExecutor::ThreadHandle*);
158   void threadStopped(
159     folly::wangle::ThreadPoolExecutor::ThreadHandle*);
160   void threadPreviouslyStarted(
161       folly::wangle::ThreadPoolExecutor::ThreadHandle* thread) {
162     threadStarted(thread);
163   }
164   void threadNotYetStopped(
165       folly::wangle::ThreadPoolExecutor::ThreadHandle* thread) {
166     threadStopped(thread);
167   }
168
169  private:
170   std::map<folly::wangle::ThreadPoolExecutor::ThreadHandle*,
171            std::shared_ptr<Acceptor>> workers_;
172   std::shared_ptr<AcceptorFactory> acceptorFactory_;
173   folly::wangle::IOThreadPoolExecutor* exec_{nullptr};
174   std::shared_ptr<std::vector<std::shared_ptr<folly::AsyncSocketBase>>> sockets_;
175   std::shared_ptr<ServerSocketFactory> socketFactory_;
176 };
177
178 template <typename F>
179 void ServerWorkerPool::forEachWorker(F&& f) const {
180   for (const auto& kv : workers_) {
181     f(kv.second.get());
182   }
183 }
184
185 class DefaultAcceptPipelineFactory
186     : public PipelineFactory<wangle::ChannelPipeline<void*, std::exception>> {
187   typedef wangle::ChannelPipeline<
188       void*,
189       std::exception> AcceptPipeline;
190
191  public:
192   AcceptPipeline* newPipeline(std::shared_ptr<AsyncSocket>) {
193     return new AcceptPipeline;
194   }
195 };
196
197 } // namespace