pub use crate::rpc::pb::etcdserverpb::alarm_request::AlarmAction;
pub use crate::rpc::pb::etcdserverpb::AlarmType;
use super::pb::etcdserverpb;
use crate::auth::AuthService;
use crate::channel::Channel;
use crate::error::Result;
use crate::rpc::pb::etcdserverpb::{
AlarmRequest as PbAlarmRequest, AlarmResponse as PbAlarmResponse,
DefragmentRequest as PbDefragmentRequest, DefragmentResponse as PbDefragmentResponse,
HashKvRequest as PbHashKvRequest, HashKvResponse as PbHashKvResponse,
HashRequest as PbHashRequest, HashResponse as PbHashResponse,
MoveLeaderRequest as PbMoveLeaderRequest, MoveLeaderResponse as PbMoveLeaderResponse,
SnapshotRequest as PbSnapshotRequest, SnapshotResponse as PbSnapshotResponse,
StatusRequest as PbStatusRequest, StatusResponse as PbStatusResponse,
};
use crate::rpc::ResponseHeader;
use etcdserverpb::maintenance_client::MaintenanceClient as PbMaintenanceClient;
use etcdserverpb::AlarmMember as PbAlarmMember;
use http::HeaderValue;
use std::sync::{Arc, RwLock};
use tonic::codec::Streaming as PbStreaming;
use tonic::{IntoRequest, Request};
#[repr(transparent)]
#[derive(Clone)]
pub struct MaintenanceClient {
inner: PbMaintenanceClient<AuthService<Channel>>,
}
#[derive(Debug, Default, Clone)]
#[repr(transparent)]
pub struct AlarmOptions(PbAlarmRequest);
impl AlarmOptions {
#[inline]
pub const fn new() -> Self {
AlarmOptions(PbAlarmRequest {
action: AlarmAction::Get as i32,
member_id: 0,
alarm: AlarmType::None as i32,
})
}
#[inline]
const fn with_action_and_type(
mut self,
alarm_action: AlarmAction,
alarm_type: AlarmType,
) -> Self {
self.0.action = alarm_action as i32;
self.0.alarm = alarm_type as i32;
self
}
#[inline]
pub fn with_member(&mut self, member: u64) {
self.0.member_id = member;
}
}
impl From<AlarmOptions> for PbAlarmRequest {
#[inline]
fn from(alarm: AlarmOptions) -> Self {
alarm.0
}
}
impl IntoRequest<PbAlarmRequest> for AlarmOptions {
#[inline]
fn into_request(self) -> Request<PbAlarmRequest> {
Request::new(self.into())
}
}
#[derive(Debug, Default, Clone)]
struct StatusOptions(PbStatusRequest);
impl StatusOptions {
#[inline]
const fn new() -> Self {
Self(PbStatusRequest {})
}
}
impl From<StatusOptions> for PbStatusRequest {
#[inline]
fn from(status: StatusOptions) -> Self {
status.0
}
}
impl IntoRequest<PbStatusRequest> for StatusOptions {
#[inline]
fn into_request(self) -> Request<PbStatusRequest> {
Request::new(self.into())
}
}
#[derive(Debug, Default, Clone)]
struct DefragmentOptions(PbDefragmentRequest);
impl DefragmentOptions {
#[inline]
const fn new() -> Self {
Self(PbDefragmentRequest {})
}
}
impl From<DefragmentOptions> for PbDefragmentRequest {
#[inline]
fn from(defragment: DefragmentOptions) -> Self {
defragment.0
}
}
impl IntoRequest<PbDefragmentRequest> for DefragmentOptions {
#[inline]
fn into_request(self) -> Request<PbDefragmentRequest> {
Request::new(self.into())
}
}
#[derive(Debug, Default, Clone)]
struct HashOptions(PbHashRequest);
impl HashOptions {
#[inline]
const fn new() -> Self {
Self(PbHashRequest {})
}
}
impl From<HashOptions> for PbHashRequest {
#[inline]
fn from(hash: HashOptions) -> Self {
hash.0
}
}
impl IntoRequest<PbHashRequest> for HashOptions {
#[inline]
fn into_request(self) -> Request<PbHashRequest> {
Request::new(self.into())
}
}
#[derive(Debug, Default, Clone)]
#[repr(transparent)]
struct HashKvOptions(PbHashKvRequest);
impl HashKvOptions {
#[inline]
const fn new(revision: i64) -> Self {
Self(PbHashKvRequest { revision })
}
}
impl From<HashKvOptions> for PbHashKvRequest {
#[inline]
fn from(hash_kv: HashKvOptions) -> Self {
hash_kv.0
}
}
impl IntoRequest<PbHashKvRequest> for HashKvOptions {
#[inline]
fn into_request(self) -> Request<PbHashKvRequest> {
Request::new(self.into())
}
}
#[derive(Debug, Default, Clone)]
struct SnapshotOptions(PbSnapshotRequest);
impl SnapshotOptions {
#[inline]
const fn new() -> Self {
Self(PbSnapshotRequest {})
}
}
impl From<SnapshotOptions> for PbSnapshotRequest {
#[inline]
fn from(snapshot: SnapshotOptions) -> Self {
snapshot.0
}
}
impl IntoRequest<PbSnapshotRequest> for SnapshotOptions {
#[inline]
fn into_request(self) -> Request<PbSnapshotRequest> {
Request::new(self.into())
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct AlarmResponse(PbAlarmResponse);
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Clone, PartialEq, Eq)]
pub struct AlarmMember {
member_id: u64,
alarm: AlarmType,
}
impl AlarmMember {
#[inline]
pub fn member_id(&self) -> u64 {
self.member_id
}
#[inline]
pub fn alarm(&self) -> AlarmType {
self.alarm
}
}
impl AlarmResponse {
#[inline]
const fn new(resp: PbAlarmResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn alarms(&self) -> &[AlarmMember] {
unsafe { &*(&self.0.alarms as *const Vec<PbAlarmMember> as *const Vec<AlarmMember>) }
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct StatusResponse(PbStatusResponse);
impl StatusResponse {
#[inline]
const fn new(resp: PbStatusResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn version(&self) -> &str {
&self.0.version
}
#[inline]
pub fn db_size(&self) -> i64 {
self.0.db_size
}
#[inline]
pub fn leader(&self) -> u64 {
self.0.leader
}
#[inline]
pub fn raft_index(&self) -> u64 {
self.0.raft_index
}
#[inline]
pub fn raft_term(&self) -> u64 {
self.0.raft_term
}
#[inline]
pub fn raft_applied_index(&self) -> u64 {
self.0.raft_applied_index
}
#[inline]
pub fn errors(&self) -> &[String] {
&self.0.errors
}
#[inline]
pub fn raft_used_db_size(&self) -> i64 {
self.0.db_size_in_use
}
#[inline]
pub fn is_learner(&self) -> bool {
self.0.is_learner
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct DefragmentResponse(PbDefragmentResponse);
impl DefragmentResponse {
#[inline]
const fn new(resp: PbDefragmentResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct HashResponse(PbHashResponse);
impl HashResponse {
#[inline]
const fn new(resp: PbHashResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn hash(&self) -> u32 {
self.0.hash
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct HashKvResponse(PbHashKvResponse);
impl HashKvResponse {
#[inline]
const fn new(resp: PbHashKvResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn hash(&self) -> u32 {
self.0.hash
}
#[inline]
pub fn compact_version(&self) -> i64 {
self.0.compact_revision
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct SnapshotResponse(PbSnapshotResponse);
impl SnapshotResponse {
#[inline]
const fn new(resp: PbSnapshotResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
#[inline]
pub fn remaining_bytes(&self) -> u64 {
self.0.remaining_bytes
}
#[inline]
pub fn blob(&self) -> &[u8] {
&self.0.blob
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug)]
#[repr(transparent)]
pub struct SnapshotStreaming(PbStreaming<PbSnapshotResponse>);
impl SnapshotStreaming {
#[inline]
pub async fn message(&mut self) -> Result<Option<SnapshotResponse>> {
let ret = self.0.message().await?;
match ret {
Some(rsp) => Ok(Some(SnapshotResponse::new(rsp))),
None => Ok(None),
}
}
}
#[derive(Debug, Default, Clone)]
#[repr(transparent)]
pub struct MoveLeaderOptions(PbMoveLeaderRequest);
impl MoveLeaderOptions {
#[inline]
const fn with_target_id(mut self, target_id: u64) -> Self {
self.0.target_id = target_id;
self
}
#[inline]
pub const fn new() -> Self {
Self(PbMoveLeaderRequest { target_id: 0 })
}
}
impl From<MoveLeaderOptions> for PbMoveLeaderRequest {
#[inline]
fn from(options: MoveLeaderOptions) -> Self {
options.0
}
}
impl IntoRequest<PbMoveLeaderRequest> for MoveLeaderOptions {
#[inline]
fn into_request(self) -> Request<PbMoveLeaderRequest> {
Request::new(self.into())
}
}
#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
#[derive(Debug, Clone)]
#[repr(transparent)]
pub struct MoveLeaderResponse(PbMoveLeaderResponse);
impl MoveLeaderResponse {
#[inline]
const fn new(resp: PbMoveLeaderResponse) -> Self {
Self(resp)
}
#[inline]
pub fn header(&self) -> Option<&ResponseHeader> {
self.0.header.as_ref().map(From::from)
}
#[inline]
pub fn take_header(&mut self) -> Option<ResponseHeader> {
self.0.header.take().map(ResponseHeader::new)
}
}
impl MaintenanceClient {
#[inline]
pub(crate) fn new(channel: Channel, auth_token: Arc<RwLock<Option<HeaderValue>>>) -> Self {
let inner = PbMaintenanceClient::new(AuthService::new(channel, auth_token));
Self { inner }
}
#[inline]
pub async fn alarm(
&mut self,
alarm_action: AlarmAction,
alarm_type: AlarmType,
options: Option<AlarmOptions>,
) -> Result<AlarmResponse> {
let resp = self
.inner
.alarm(
options
.unwrap_or_default()
.with_action_and_type(alarm_action, alarm_type),
)
.await?
.into_inner();
Ok(AlarmResponse::new(resp))
}
#[inline]
pub async fn status(&mut self) -> Result<StatusResponse> {
let resp = self.inner.status(StatusOptions::new()).await?.into_inner();
Ok(StatusResponse::new(resp))
}
#[inline]
pub async fn defragment(&mut self) -> Result<DefragmentResponse> {
let resp = self
.inner
.defragment(DefragmentOptions::new())
.await?
.into_inner();
Ok(DefragmentResponse::new(resp))
}
#[inline]
pub async fn hash(&mut self) -> Result<HashResponse> {
let resp = self.inner.hash(HashOptions::new()).await?.into_inner();
Ok(HashResponse::new(resp))
}
#[inline]
pub async fn hash_kv(&mut self, revision: i64) -> Result<HashKvResponse> {
let resp = self
.inner
.hash_kv(HashKvOptions::new(revision))
.await?
.into_inner();
Ok(HashKvResponse::new(resp))
}
#[inline]
pub async fn snapshot(&mut self) -> Result<SnapshotStreaming> {
let resp = self
.inner
.snapshot(SnapshotOptions::new())
.await?
.into_inner();
Ok(SnapshotStreaming(resp))
}
#[inline]
pub async fn move_leader(&mut self, target_id: u64) -> Result<MoveLeaderResponse> {
let resp = self
.inner
.move_leader(MoveLeaderOptions::new().with_target_id(target_id))
.await?
.into_inner();
Ok(MoveLeaderResponse::new(resp))
}
}