Get down to a single write call

buffer.cpp

@@ -44,6 +44,15 @@ size_t Buffer::WriteMaxLen() const {
 	return size_ - len_;
 }
 
+bool Buffer::Write(const std::string_view& str) {
+	if (WriteMaxLen() < str.size()) {
+		return false;
+	}
+	memcpy(WritePtr(), str.data(), str.size());
+	Wrote(str.size());
+	return true;
+}
+
 void Buffer::Wrote(size_t len) {
 	CHECK_LE(len, WriteMaxLen());
 	len_ += len;

buffer.h

@@ -26,6 +26,7 @@ class Buffer : public ConstBuffer {
 
 	[[nodiscard]] char *WritePtr();
 	[[nodiscard]] size_t WriteMaxLen() const;
+	bool Write(const std::string_view& str);
 	void Wrote(size_t len);
 
 	void Commit(); // commit read position

example_simple.cpp

@@ -10,8 +10,9 @@ int main(int argc, char *argv[]) {
 	gflags::ParseCommandLineFlags(&argc, &argv, true);
 
 	FastCGIServer server(FLAGS_port, [](std::unique_ptr<FastCGIRequest> request) {
-		request->Write({{"Content-Type", "text/plain"}}, {"Hello world"});
-		request->WriteEnd();
+		request->WriteHeader("Content-Type", "text/plain");
+		request->WriteBody("Hello world");
+		request->End();
 	});
 	server.Serve();
 }

fastcgi_conn.cpp

@@ -22,44 +22,12 @@ FastCGIConn::~FastCGIConn() {
 	LOG(INFO) << "connection closed";
 }
 
-void FastCGIConn::WriteBlock(uint8_t type, uint16_t request_id, const std::vector<iovec>& vecs) {
-	std::vector<iovec> out_vecs;
-	out_vecs.reserve(vecs.size() + 1);
-
-	FastCGIHeader header;
-	header.version = 1;
-	header.type = type;
-	header.SetRequestId(request_id);
-	uint16_t content_length = 0;
-	for (auto& vec : vecs) {
-		content_length += vec.iov_len;
+void FastCGIConn::Write(const std::vector<iovec>& vecs) {
+	size_t total_size = 0;
+	for (const auto& vec : vecs) {
+		total_size += vec.iov_len;
 	}
-	header.SetContentLength(content_length);
-	out_vecs.push_back(iovec{
-		.iov_base = &header,
-		.iov_len = sizeof(header),
-	});
-
-	for (auto& vec : vecs) {
-		out_vecs.push_back(vec);
-	}
-
-	CHECK_EQ(writev(sock_, out_vecs.data(), out_vecs.size()), content_length + sizeof(header));
-}
-
-void FastCGIConn::WriteOutput(uint16_t request_id, const std::vector<iovec>& vecs) {
-	WriteBlock(6, request_id, vecs);
-}
-
-void FastCGIConn::WriteEnd(uint16_t request_id) {
-	FastCGIEndRequest end;
-
-	std::vector<iovec> vecs;
-	vecs.push_back(iovec{
-		.iov_base = &end,
-		.iov_len = sizeof(end),
-	});
-	WriteBlock(3, request_id, vecs);
+	CHECK_EQ(writev(sock_, vecs.data(), vecs.size()), total_size);
 }
 
 void FastCGIConn::Serve() {

fastcgi_conn.h

@@ -15,9 +15,7 @@ class FastCGIConn {
 
 	void Serve();
 
-	void WriteBlock(uint8_t type, uint16_t request_id, const std::vector<iovec>& vecs);
-	void WriteOutput(uint16_t request_id, const std::vector<iovec>& vecs);
-	void WriteEnd(uint16_t request_id);
+	void Write(const std::vector<iovec>& vecs);
 
   private:
   	const int sock_;

fastcgi_parse.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <arpa/inet.h>
+
 struct FastCGIHeader {
 	uint8_t version;
 	uint8_t type;

fastcgi_request.cpp

@@ -1,12 +1,14 @@
 #include <sys/uio.h>
 
-#include "fastcgi_conn.h"
-
 #include "fastcgi_request.h"
 
+#include "fastcgi_conn.h"
+#include "fastcgi_parse.h"
+
 FastCGIRequest::FastCGIRequest(uint16_t request_id, FastCGIConn* conn)
 		: request_id_(request_id),
-		  conn_(conn) {}
+		  conn_(conn),
+		  out_buf_(fastcgi_max_record_len) {}
 
 void FastCGIRequest::AddParam(const std::string_view& key, const std::string_view& value) {
 	params_.try_emplace(std::string(key), std::string(value));
@@ -20,49 +22,57 @@ const std::string& FastCGIRequest::GetParam(const std::string& key) {
 	return params_.at(key);
 }
 
-void FastCGIRequest::Write(const std::vector<std::pair<std::string_view, std::string_view>>& headers, const std::vector<std::string_view>& body) {
-	std::vector<iovec> vecs;
-	vecs.reserve((headers.size() * 4) + 1 + body.size());
+void FastCGIRequest::WriteHeader(const std::string_view& name, const std::string_view& value) {
+	CHECK(!body_sent_);
+	CHECK(out_buf_.Write(name));
+	CHECK(out_buf_.Write(": "));
+	CHECK(out_buf_.Write(value));
+	CHECK(out_buf_.Write("\n"));
+}
 
-	CHECK(headers.empty() || !body_sent_);
-
-	for (const auto& header : headers) {
-		vecs.push_back(iovec{
-			.iov_base = const_cast<char*>(header.first.data()),
-			.iov_len = header.first.size(),
-		});
-		vecs.push_back(iovec{
-			.iov_base = const_cast<char*>(": "),
-			.iov_len = 2,
-		});
-		vecs.push_back(iovec{
-			.iov_base = const_cast<char*>(header.second.data()),
-			.iov_len = header.second.size(),
-		});
-		vecs.push_back(iovec{
-			.iov_base = const_cast<char*>("\n"),
-			.iov_len = 1,
-		});
-	}
-
-	if (!body.empty() && !body_sent_) {
+void FastCGIRequest::WriteBody(const std::string_view& body) {
+	if (!body_sent_) {
+		CHECK(out_buf_.Write("\n"));
 		body_sent_ = true;
-		vecs.push_back(iovec{
-			.iov_base = const_cast<char*>("\n"),
-			.iov_len = 1,
-		});
 	}
-
-	for (const auto& chunk : body) {
-		vecs.push_back(iovec{
-			.iov_base = const_cast<char*>(chunk.data()),
-			.iov_len = chunk.size(),
-		});
-	}
-
-	conn_->WriteOutput(request_id_, vecs);
+	// TODO: make this able to span multiple packets
+	CHECK(out_buf_.Write(body));
 }
 
-void FastCGIRequest::WriteEnd() {
-	conn_->WriteEnd(request_id_);
+void FastCGIRequest::End() {
+	FastCGIHeader output_header;
+	FastCGIHeader end_header;
+	FastCGIEndRequest end;
+
+	const auto output_len = out_buf_.ReadMaxLen();
+
+	output_header.version = 1;
+	output_header.type = 6;
+	output_header.SetRequestId(request_id_);
+	output_header.SetContentLength(output_len);
+
+	end_header.version = 1;
+	end_header.type = 3;
+	end_header.SetRequestId(request_id_);
+	end_header.SetContentLength(sizeof(end));
+
+	std::vector<iovec> vecs{
+		iovec{
+			.iov_base = &output_header,
+			.iov_len = sizeof(output_header),
+		},
+		iovec{
+			.iov_base = (void *)(CHECK_NOTNULL(out_buf_.Read(output_len))),
+			.iov_len = output_len,
+		},
+		{
+			.iov_base = &end_header,
+			.iov_len = sizeof(end_header),
+		},
+		{
+			.iov_base = &end,
+			.iov_len = sizeof(end),
+		},
+	};
+	conn_->Write(vecs);
 }

fastcgi_request.h

@@ -2,6 +2,8 @@
 
 #include <unordered_map>
 
+#include "buffer.h"
+
 class FastCGIConn;
 
 class FastCGIRequest {
@@ -13,8 +15,10 @@ class FastCGIRequest {
 
 	const std::string& GetParam(const std::string& key);
 
-	void Write(const std::vector<std::pair<std::string_view, std::string_view>>& headers, const std::vector<std::string_view>& body);
-	void WriteEnd();
+	void WriteHeader(const std::string_view& name, const std::string_view& value);
+	void WriteBody(const std::string_view& body);
+	void Flush();
+	void End();
 
   private:
 	const uint16_t request_id_;
@@ -22,5 +26,7 @@ class FastCGIRequest {
 
 	std::unordered_map<std::string, std::string> params_;
 	std::string in_;
+
+	Buffer out_buf_;
 	bool body_sent_ = false;
 };