session.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef TVM_RUNTIME_CRT_RPC_COMMON_SESSION_H_
#define TVM_RUNTIME_CRT_RPC_COMMON_SESSION_H_

#include <inttypes.h>
#include <tvm/runtime/crt/error_codes.h>
#include <tvm/runtime/crt/rpc_common/frame_buffer.h>
#include <tvm/runtime/crt/rpc_common/framing.h>
#include <tvm/runtime/crt/rpc_common/write_stream.h>

namespace tvm {
namespace runtime {
namespace micro_rpc {

enum class MessageType : uint8_t {
  kStartSessionInit = 0x00,
  kStartSessionReply = 0x01,
  kTerminateSession = 0x02,
  kLog = 0x03,
  kNormal = 0x10,
};

#if defined(_MSC_VER)

#pragma pack(push, 1)
typedef struct SessionHeader {
  uint16_t session_id;
  MessageType message_type;
} SessionHeader;
#pragma pack(pop)

#else

typedef struct SessionHeader {
  uint16_t session_id;
  MessageType message_type;
} __attribute__((packed)) SessionHeader;

#endif

class Session {
 public:
  typedef void (*MessageReceivedFunc)(void* context, MessageType message_type, FrameBuffer* buf);

  static constexpr const uint8_t kInvalidNonce = 0;

  Session(Framer* framer, FrameBuffer* receive_buffer, MessageReceivedFunc message_received_func,
          void* message_received_func_context)
      : local_nonce_{kInvalidNonce},
        session_id_{0},
        state_{State::kReset},
        receiver_{this},
        framer_{framer},
        receive_buffer_{receive_buffer},
        receive_buffer_has_complete_message_{false},
        message_received_func_{message_received_func},
        message_received_func_context_{message_received_func_context} {
    // Session can be used for system startup logging, before the RPC server is instantiated. In
    // this case, allow receive_buffer_ to be nullptr. The instantiator agrees not to use
    // Receiver().
    if (receive_buffer_ != nullptr) {
      receive_buffer_->Clear();
    }
  }

  tvm_crt_error_t Initialize(uint8_t initial_session_nonce);

  tvm_crt_error_t TerminateSession();

  tvm_crt_error_t StartSession();

  WriteStream* Receiver() { return &receiver_; }

  tvm_crt_error_t SendMessage(MessageType message_type, const uint8_t* message_data,
                              size_t message_size_bytes);

  tvm_crt_error_t StartMessage(MessageType message_type, size_t message_size_bytes);

  tvm_crt_error_t SendBodyChunk(const uint8_t* chunk_data, size_t chunk_size_bytes);

  tvm_crt_error_t FinishMessage();

  bool IsEstablished() const { return state_ == State::kSessionEstablished; }

  void ClearReceiveBuffer();

  static const constexpr uint8_t kVersion = 0x01;

 private:
  class SessionReceiver : public WriteStream {
   public:
    explicit SessionReceiver(Session* session) : session_{session} {}
    virtual ~SessionReceiver() {}

    ssize_t Write(const uint8_t* data, size_t data_size_bytes) override;
    void PacketDone(bool is_valid) override;

   private:
    void operator delete(void*) noexcept {}  // NOLINT(readability/casting)
    Session* session_;
  };

  enum class State : uint8_t {
    kReset = 0,
    kNoSessionEstablished = 1,
    kStartSessionSent = 2,
    kSessionEstablished = 3,
  };

  void RegenerateNonce();

  tvm_crt_error_t SendInternal(MessageType message_type, const uint8_t* message_data,
                               size_t message_size_bytes);

  void SendSessionStartReply(const SessionHeader& header);

  void ProcessStartSessionInit(const SessionHeader& header);

  void ProcessStartSessionReply(const SessionHeader& header);

  void OnSessionEstablishedMessage();

  void OnSessionTerminatedMessage();

  void SetSessionId(uint8_t initiator_nonce, uint8_t responder_nonce) {
    session_id_ = initiator_nonce | (((uint16_t)responder_nonce) << 8);
  }

  uint8_t InitiatorNonce(uint16_t session_id) { return session_id & 0xff; }

  uint8_t ResponderNonce(uint16_t session_id) { return (session_id >> 8) & 0xff; }

  uint8_t local_nonce_;
  uint16_t session_id_;
  State state_;
  SessionReceiver receiver_;
  Framer* framer_;
  FrameBuffer* receive_buffer_;
  bool receive_buffer_has_complete_message_;
  MessageReceivedFunc message_received_func_;
  void* message_received_func_context_;
};

}  // namespace micro_rpc
}  // namespace runtime
}  // namespace tvm

#endif  // TVM_RUNTIME_CRT_RPC_COMMON_SESSION_H_