largefileclient.cpp

#include <grpc/grpc.h>
#include <grpcpp/ext/proto_server_reflection_plugin.h>
#include <grpcpp/grpcpp.h>

#include <cstddef>
#include <fstream>
#include <ostream>
#include <string>

#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/log/initialize.h"

#include "largefiles.grpc.pb.h"

ABSL_FLAG(std::string, target, "localhost:50051", "Server address");
ABSL_FLAG(std::string, file, "", "File to upload");
ABSL_FLAG(size_t, buffer, 16 * 1024, "Size of file chunks to stream");

using grpc::CallbackServerContext;
using grpc::Channel;
using grpc::ClientContext;
using grpc::Server;
using grpc::ServerBuilder;
using grpc::ServerUnaryReactor;
using grpc::Status;
using largefiles::Chunk;
using largefiles::FileStreamBlock;
using largefiles::FileUploadResponse;
using largefiles::FileWriter;

namespace
{

    // Sends requests as quickly as possible and times how long it takes to perform
    // the write operation.
    class UploadFileClientReactor final
        : public grpc::ClientWriteReactor<largefiles::FileStreamBlock>
    {
    public:
        explicit UploadFileClientReactor(std::string fileToUpload, size_t bufferSize) 
        : fileToUpload_(fileToUpload)
        , bufferSize_(bufferSize) {}

        void Start()
        {
            absl::MutexLock lock(&mu_);
            StartCall();
            Write();
        }

        void awaitDone() {
            absl::MutexLock lock(&mu_);
            mu_.Await(absl::Condition(+[](bool *done)
                                      { return *done; },
                                      &done_));
        }

        ~UploadFileClientReactor() override
        {
           awaitDone();
           fileStream_.close();
        }

        void OnWriteDone(bool ok) override
        {
            absl::MutexLock lock(&mu_);
            std::cout << "Writing took " << absl::Now() - *time_ << std::endl;
            time_ = std::nullopt;
            if (ok)
            {
                Write();
            }
        }

        void OnDone(const grpc::Status &status) override
        {
            if (status.ok())
            {
                std::cout << "Done\n";
            }
            else
            {
                std::cout << "Done with error: [" << status.error_code() << "] "
                          << status.error_message() << "\n";
            }
            absl::MutexLock lock(&mu_);
            done_ = true;
        }

    private:
        void Write() ABSL_EXCLUSIVE_LOCKS_REQUIRED(&mu_)
        {
            if (reqs_ == 0) {
                req_.mutable_metadata()->set_name(fileToUpload_);
                req_.mutable_metadata()->set_populated(true);

                // Open file on first write
                fileStream_.open(fileToUpload_, std::ios::in | std::ios::binary);
                if (!fileStream_) {
                    std::cerr << "Failed to open file: " << fileToUpload_ << std::endl;
                    throw std::runtime_error("Failed to open file" + fileToUpload_);
                }
            } else {
                req_.clear_metadata();
            }

            // Write next chunk
            std::vector<char> buffer(bufferSize_);
            if (fileStream_.read(buffer.data(), bufferSize_) || fileStream_.gcount() > 0) {
                req_.mutable_filedata()->mutable_chunk()->assign(buffer.data(), fileStream_.gcount());
                std::cout << "Writing chunk of size " << fileStream_.gcount() << std::endl;
                reqs_++;
                StartWrite(&req_);
                time_ = absl::Now();
            } else {
                std::cout << "Closing file. All data sent." << std::endl;
                fileStream_.close();
                StartWritesDone();
            }

        }

        std::string fileToUpload_;
        size_t bufferSize_;
        absl::Mutex mu_;
        bool done_ ABSL_GUARDED_BY(&mu_) = false;
        largefiles::FileStreamBlock req_;
        size_t reqs_;
        std::optional<absl::Time> time_ ABSL_GUARDED_BY(mu_);
        std::ifstream fileStream_;
    };

} // namespace

int main(int argc, char **argv)
{
    absl::ParseCommandLine(argc, argv);
    absl::InitializeLog();

    if (absl::GetFlag(FLAGS_file).empty())
    {
        std::cerr << "Please specify a file to upload with --file=<filename>\n";
        return 1;
    }
    grpc::ChannelArguments channel_arguments;
    auto channel = grpc::CreateCustomChannel(absl::GetFlag(FLAGS_target),
                                             grpc::InsecureChannelCredentials(),
                                             channel_arguments);
    auto stub = FileWriter::NewStub(channel);
    UploadFileClientReactor reactor(absl::GetFlag(FLAGS_file), absl::GetFlag(FLAGS_buffer));
    grpc::ClientContext context;
    ::largefiles::FileUploadResponse response;
    stub->async()->UploadFile(&context, &response, &reactor);
    reactor.Start();
    reactor.awaitDone();

    std::cout << "File upload completed with response size: " << response.readsize() << std::endl;

    return 0;
}