kevo/pkg/client/client.go
2025-05-17 14:58:26 -06:00

822 lines
24 KiB
Go

package client
import (
"context"
"encoding/json"
"errors"
"fmt"
"sync"
"time"
"github.com/KevoDB/kevo/pkg/transport"
"google.golang.org/grpc/keepalive"
)
// CompressionType represents a compression algorithm
type CompressionType = transport.CompressionType
// Compression options
const (
CompressionNone = transport.CompressionNone
CompressionGzip = transport.CompressionGzip
CompressionSnappy = transport.CompressionSnappy
)
// ClientOptions configures a Kevo client
type ClientOptions struct {
// Connection options
Endpoint string // Server address
ConnectTimeout time.Duration // Timeout for connection attempts
RequestTimeout time.Duration // Default timeout for requests
TransportType string // Transport type (e.g. "grpc")
PoolSize int // Connection pool size
// Security options
TLSEnabled bool // Enable TLS
CertFile string // Client certificate file
KeyFile string // Client key file
CAFile string // CA certificate file
// Retry options
MaxRetries int // Maximum number of retries
InitialBackoff time.Duration // Initial retry backoff
MaxBackoff time.Duration // Maximum retry backoff
BackoffFactor float64 // Backoff multiplier
RetryJitter float64 // Random jitter factor
// Performance options
Compression CompressionType // Compression algorithm
MaxMessageSize int // Maximum message size
// Keepalive options
KeepaliveTime time.Duration // Time between keepalive pings (0 for default)
KeepaliveTimeout time.Duration // Time to wait for ping ack (0 for default)
}
// DefaultClientOptions returns sensible default client options
func DefaultClientOptions() ClientOptions {
return ClientOptions{
Endpoint: "localhost:50051",
ConnectTimeout: time.Second * 5,
RequestTimeout: time.Second * 10,
TransportType: "grpc",
PoolSize: 5,
TLSEnabled: false,
MaxRetries: 3,
InitialBackoff: time.Millisecond * 100,
MaxBackoff: time.Second * 2,
BackoffFactor: 1.5,
RetryJitter: 0.2,
Compression: CompressionNone,
MaxMessageSize: 16 * 1024 * 1024, // 16MB
KeepaliveTime: 30 * time.Second, // 30 seconds keepalive time
KeepaliveTimeout: 10 * time.Second, // 10 seconds timeout
}
}
// ReplicaInfo represents information about a replica node
type ReplicaInfo struct {
Address string // Host:port of the replica
LastSequence uint64 // Last applied sequence number
Available bool // Whether the replica is available
Region string // Optional region information
Meta map[string]string // Additional metadata
}
// NodeInfo contains information about the server node and topology
type NodeInfo struct {
Role string // "primary", "replica", or "standalone"
PrimaryAddr string // Address of the primary node
Replicas []ReplicaInfo // Available replica nodes
LastSequence uint64 // Last applied sequence number
ReadOnly bool // Whether the node is in read-only mode
}
// Client represents a connection to a Kevo database server
type Client struct {
options ClientOptions
client transport.Client
primaryConn transport.Client // Connection to primary (when connected to replica)
replicaConn []transport.Client // Connections to replicas (when connected to primary)
nodeInfo *NodeInfo // Information about the current node and topology
connMutex sync.RWMutex // Protects connections
}
// NewClient creates a new Kevo client with the given options
func NewClient(options ClientOptions) (*Client, error) {
if options.Endpoint == "" {
return nil, errors.New("endpoint is required")
}
// Configure keepalive parameters if specified
var keepaliveParams *keepalive.ClientParameters
if options.KeepaliveTime > 0 && options.KeepaliveTimeout > 0 {
keepaliveParams = &keepalive.ClientParameters{
Time: options.KeepaliveTime,
Timeout: options.KeepaliveTimeout,
PermitWithoutStream: true, // Allow pings even when there are no active streams
}
}
transportOpts := transport.TransportOptions{
Timeout: options.ConnectTimeout,
MaxMessageSize: options.MaxMessageSize,
Compression: options.Compression,
TLSEnabled: options.TLSEnabled,
CertFile: options.CertFile,
KeyFile: options.KeyFile,
CAFile: options.CAFile,
KeepaliveParams: keepaliveParams,
RetryPolicy: transport.RetryPolicy{
MaxRetries: options.MaxRetries,
InitialBackoff: options.InitialBackoff,
MaxBackoff: options.MaxBackoff,
BackoffFactor: options.BackoffFactor,
Jitter: options.RetryJitter,
},
}
transportClient, err := transport.GetClient(options.TransportType, options.Endpoint, transportOpts)
if err != nil {
return nil, fmt.Errorf("failed to create transport client: %w", err)
}
return &Client{
options: options,
client: transportClient,
}, nil
}
// Connect establishes a connection to the server
// and discovers the replication topology if available
func (c *Client) Connect(ctx context.Context) error {
// First connect to the primary endpoint
if err := c.client.Connect(ctx); err != nil {
return err
}
// Query node information to discover the topology
return c.discoverTopology(ctx)
}
// discoverTopology queries the node for replication information
// and establishes additional connections if needed
func (c *Client) discoverTopology(ctx context.Context) error {
// Get node info from the connected server
nodeInfo, err := c.getNodeInfo(ctx)
if err != nil {
// If GetNodeInfo isn't supported, assume it's standalone
// This ensures backward compatibility with older servers
nodeInfo = &NodeInfo{
Role: "standalone",
ReadOnly: false,
}
}
c.connMutex.Lock()
defer c.connMutex.Unlock()
// Store the node info
c.nodeInfo = nodeInfo
// Based on the role, establish additional connections as needed
switch nodeInfo.Role {
case "replica":
// If connected to a replica and a primary is available, connect to it
if nodeInfo.PrimaryAddr != "" && nodeInfo.PrimaryAddr != c.options.Endpoint {
primaryOptions := c.options
primaryOptions.Endpoint = nodeInfo.PrimaryAddr
// Create client connection to primary
primaryClient, err := transport.GetClient(
primaryOptions.TransportType,
primaryOptions.Endpoint,
c.createTransportOptions(primaryOptions),
)
if err == nil {
// Try to connect to primary
if err := primaryClient.Connect(ctx); err == nil {
c.primaryConn = primaryClient
}
}
}
case "primary":
// If connected to a primary and replicas are available, connect to some of them
c.replicaConn = make([]transport.Client, 0, len(nodeInfo.Replicas))
// Connect to up to 2 replicas (to avoid too many connections)
for i, replica := range nodeInfo.Replicas {
if i >= 2 || !replica.Available {
continue
}
replicaOptions := c.options
replicaOptions.Endpoint = replica.Address
// Create client connection to replica
replicaClient, err := transport.GetClient(
replicaOptions.TransportType,
replicaOptions.Endpoint,
c.createTransportOptions(replicaOptions),
)
if err == nil {
// Try to connect to replica
if err := replicaClient.Connect(ctx); err == nil {
c.replicaConn = append(c.replicaConn, replicaClient)
}
}
}
}
return nil
}
// createTransportOptions converts client options to transport options
func (c *Client) createTransportOptions(options ClientOptions) transport.TransportOptions {
// Configure keepalive parameters if specified
var keepaliveParams *keepalive.ClientParameters
if options.KeepaliveTime > 0 && options.KeepaliveTimeout > 0 {
keepaliveParams = &keepalive.ClientParameters{
Time: options.KeepaliveTime,
Timeout: options.KeepaliveTimeout,
PermitWithoutStream: true, // Allow pings even when there are no active streams
}
}
return transport.TransportOptions{
Timeout: options.ConnectTimeout,
MaxMessageSize: options.MaxMessageSize,
Compression: options.Compression,
TLSEnabled: options.TLSEnabled,
CertFile: options.CertFile,
KeyFile: options.KeyFile,
CAFile: options.CAFile,
KeepaliveParams: keepaliveParams,
RetryPolicy: transport.RetryPolicy{
MaxRetries: options.MaxRetries,
InitialBackoff: options.InitialBackoff,
MaxBackoff: options.MaxBackoff,
BackoffFactor: options.BackoffFactor,
Jitter: options.RetryJitter,
},
}
}
// Close closes all connections to servers
func (c *Client) Close() error {
c.connMutex.Lock()
defer c.connMutex.Unlock()
// Close primary connection
if c.primaryConn != nil {
c.primaryConn.Close()
c.primaryConn = nil
}
// Close replica connections
for _, replica := range c.replicaConn {
replica.Close()
}
c.replicaConn = nil
// Close main connection
return c.client.Close()
}
// getNodeInfo retrieves node information from the server
func (c *Client) getNodeInfo(ctx context.Context) (*NodeInfo, error) {
// Create a request to the GetNodeInfo endpoint
req := transport.NewRequest("GetNodeInfo", nil)
// Send the request
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
resp, err := c.client.Send(timeoutCtx, req)
if err != nil {
return nil, fmt.Errorf("failed to get node info: %w", err)
}
// Parse the response
var nodeInfoResp struct {
NodeRole int `json:"node_role"`
PrimaryAddress string `json:"primary_address"`
Replicas []json.RawMessage `json:"replicas"`
LastSequence uint64 `json:"last_sequence"`
ReadOnly bool `json:"read_only"`
}
if err := json.Unmarshal(resp.Payload(), &nodeInfoResp); err != nil {
return nil, fmt.Errorf("failed to unmarshal node info response: %w", err)
}
// Convert role from int to string
var role string
switch nodeInfoResp.NodeRole {
case 0:
role = "standalone"
case 1:
role = "primary"
case 2:
role = "replica"
default:
role = "unknown"
}
// Parse replica information
replicas := make([]ReplicaInfo, 0, len(nodeInfoResp.Replicas))
for _, rawReplica := range nodeInfoResp.Replicas {
var replica struct {
Address string `json:"address"`
LastSequence uint64 `json:"last_sequence"`
Available bool `json:"available"`
Region string `json:"region"`
Meta map[string]string `json:"meta"`
}
if err := json.Unmarshal(rawReplica, &replica); err != nil {
continue // Skip replicas that can't be parsed
}
replicas = append(replicas, ReplicaInfo{
Address: replica.Address,
LastSequence: replica.LastSequence,
Available: replica.Available,
Region: replica.Region,
Meta: replica.Meta,
})
}
return &NodeInfo{
Role: role,
PrimaryAddr: nodeInfoResp.PrimaryAddress,
Replicas: replicas,
LastSequence: nodeInfoResp.LastSequence,
ReadOnly: nodeInfoResp.ReadOnly,
}, nil
}
// IsConnected returns whether the client is connected to the server
func (c *Client) IsConnected() bool {
return c.client != nil && c.client.IsConnected()
}
// Get retrieves a value by key
// If connected to a primary with replicas, it will route reads to a replica
func (c *Client) Get(ctx context.Context, key []byte) ([]byte, bool, error) {
if !c.IsConnected() {
return nil, false, errors.New("not connected to server")
}
// Check if we should route to replica
c.connMutex.RLock()
shouldUseReplica := c.nodeInfo != nil &&
c.nodeInfo.Role == "primary" &&
len(c.replicaConn) > 0
c.connMutex.RUnlock()
req := struct {
Key []byte `json:"key"`
}{
Key: key,
}
reqData, err := json.Marshal(req)
if err != nil {
return nil, false, fmt.Errorf("failed to marshal request: %w", err)
}
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
var resp transport.Response
var sendErr error
if shouldUseReplica {
// Select a replica for reading
c.connMutex.RLock()
selectedReplica := c.replicaConn[0] // Simple selection: always use first replica
c.connMutex.RUnlock()
// Try the replica first
resp, sendErr = selectedReplica.Send(timeoutCtx, transport.NewRequest(transport.TypeGet, reqData))
// If replica fails, fall back to primary
if sendErr != nil {
resp, sendErr = c.client.Send(timeoutCtx, transport.NewRequest(transport.TypeGet, reqData))
}
} else {
// Use default connection
resp, sendErr = c.client.Send(timeoutCtx, transport.NewRequest(transport.TypeGet, reqData))
}
if sendErr != nil {
return nil, false, fmt.Errorf("failed to send request: %w", sendErr)
}
var getResp struct {
Value []byte `json:"value"`
Found bool `json:"found"`
}
if err := json.Unmarshal(resp.Payload(), &getResp); err != nil {
return nil, false, fmt.Errorf("failed to unmarshal response: %w", err)
}
return getResp.Value, getResp.Found, nil
}
// Put stores a key-value pair
// If connected to a replica, it will automatically route the write to the primary
func (c *Client) Put(ctx context.Context, key, value []byte, sync bool) (bool, error) {
if !c.IsConnected() {
return false, errors.New("not connected to server")
}
// Check if we should route to primary
c.connMutex.RLock()
shouldUsePrimary := c.nodeInfo != nil &&
c.nodeInfo.Role == "replica" &&
c.primaryConn != nil
c.connMutex.RUnlock()
req := struct {
Key []byte `json:"key"`
Value []byte `json:"value"`
Sync bool `json:"sync"`
}{
Key: key,
Value: value,
Sync: sync,
}
reqData, err := json.Marshal(req)
if err != nil {
return false, fmt.Errorf("failed to marshal request: %w", err)
}
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
var resp transport.Response
var sendErr error
if shouldUsePrimary {
// Use primary connection for writes when connected to replica
c.connMutex.RLock()
primaryConn := c.primaryConn
c.connMutex.RUnlock()
resp, sendErr = primaryConn.Send(timeoutCtx, transport.NewRequest(transport.TypePut, reqData))
} else {
// Use default connection
resp, sendErr = c.client.Send(timeoutCtx, transport.NewRequest(transport.TypePut, reqData))
// If we get a read-only error and we have node info, try to extract primary address
if sendErr != nil && c.nodeInfo == nil {
// Try to discover topology to get primary address
if discoverErr := c.discoverTopology(ctx); discoverErr == nil {
// Check again if we now have a primary connection
c.connMutex.RLock()
primaryAvailable := c.nodeInfo != nil &&
c.nodeInfo.Role == "replica" &&
c.primaryConn != nil
primaryConn := c.primaryConn
c.connMutex.RUnlock()
// If we now have a primary connection, retry the write
if primaryAvailable && primaryConn != nil {
resp, sendErr = primaryConn.Send(timeoutCtx, transport.NewRequest(transport.TypePut, reqData))
}
}
}
}
if sendErr != nil {
return false, fmt.Errorf("failed to send request: %w", sendErr)
}
var putResp struct {
Success bool `json:"success"`
}
if err := json.Unmarshal(resp.Payload(), &putResp); err != nil {
return false, fmt.Errorf("failed to unmarshal response: %w", err)
}
return putResp.Success, nil
}
// Delete removes a key-value pair
// If connected to a replica, it will automatically route the delete to the primary
func (c *Client) Delete(ctx context.Context, key []byte, sync bool) (bool, error) {
if !c.IsConnected() {
return false, errors.New("not connected to server")
}
// Check if we should route to primary
c.connMutex.RLock()
shouldUsePrimary := c.nodeInfo != nil &&
c.nodeInfo.Role == "replica" &&
c.primaryConn != nil
c.connMutex.RUnlock()
req := struct {
Key []byte `json:"key"`
Sync bool `json:"sync"`
}{
Key: key,
Sync: sync,
}
reqData, err := json.Marshal(req)
if err != nil {
return false, fmt.Errorf("failed to marshal request: %w", err)
}
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
var resp transport.Response
var sendErr error
if shouldUsePrimary {
// Use primary connection for writes when connected to replica
c.connMutex.RLock()
primaryConn := c.primaryConn
c.connMutex.RUnlock()
resp, sendErr = primaryConn.Send(timeoutCtx, transport.NewRequest(transport.TypeDelete, reqData))
} else {
// Use default connection
resp, sendErr = c.client.Send(timeoutCtx, transport.NewRequest(transport.TypeDelete, reqData))
// If we get a read-only error and we have node info, try to extract primary address
if sendErr != nil && c.nodeInfo == nil {
// Try to discover topology to get primary address
if discoverErr := c.discoverTopology(ctx); discoverErr == nil {
// Check again if we now have a primary connection
c.connMutex.RLock()
primaryAvailable := c.nodeInfo != nil &&
c.nodeInfo.Role == "replica" &&
c.primaryConn != nil
primaryConn := c.primaryConn
c.connMutex.RUnlock()
// If we now have a primary connection, retry the delete
if primaryAvailable && primaryConn != nil {
resp, sendErr = primaryConn.Send(timeoutCtx, transport.NewRequest(transport.TypeDelete, reqData))
}
}
}
}
if sendErr != nil {
return false, fmt.Errorf("failed to send request: %w", sendErr)
}
var deleteResp struct {
Success bool `json:"success"`
}
if err := json.Unmarshal(resp.Payload(), &deleteResp); err != nil {
return false, fmt.Errorf("failed to unmarshal response: %w", err)
}
return deleteResp.Success, nil
}
// BatchOperation represents a single operation in a batch
type BatchOperation struct {
Type string // "put" or "delete"
Key []byte
Value []byte // only used for "put" operations
}
// BatchWrite performs multiple operations in a single atomic batch
// If connected to a replica, it will automatically route the batch to the primary
func (c *Client) BatchWrite(ctx context.Context, operations []BatchOperation, sync bool) (bool, error) {
if !c.IsConnected() {
return false, errors.New("not connected to server")
}
// Check if we should route to primary
c.connMutex.RLock()
shouldUsePrimary := c.nodeInfo != nil &&
c.nodeInfo.Role == "replica" &&
c.primaryConn != nil
c.connMutex.RUnlock()
req := struct {
Operations []struct {
Type string `json:"type"`
Key []byte `json:"key"`
Value []byte `json:"value"`
} `json:"operations"`
Sync bool `json:"sync"`
}{
Sync: sync,
}
for _, op := range operations {
req.Operations = append(req.Operations, struct {
Type string `json:"type"`
Key []byte `json:"key"`
Value []byte `json:"value"`
}{
Type: op.Type,
Key: op.Key,
Value: op.Value,
})
}
reqData, err := json.Marshal(req)
if err != nil {
return false, fmt.Errorf("failed to marshal request: %w", err)
}
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
var resp transport.Response
var sendErr error
if shouldUsePrimary {
// Use primary connection for writes when connected to replica
c.connMutex.RLock()
primaryConn := c.primaryConn
c.connMutex.RUnlock()
resp, sendErr = primaryConn.Send(timeoutCtx, transport.NewRequest(transport.TypeBatchWrite, reqData))
} else {
// Use default connection
resp, sendErr = c.client.Send(timeoutCtx, transport.NewRequest(transport.TypeBatchWrite, reqData))
// If we get a read-only error and we have node info, try to extract primary address
if sendErr != nil && c.nodeInfo == nil {
// Try to discover topology to get primary address
if discoverErr := c.discoverTopology(ctx); discoverErr == nil {
// Check again if we now have a primary connection
c.connMutex.RLock()
primaryAvailable := c.nodeInfo != nil &&
c.nodeInfo.Role == "replica" &&
c.primaryConn != nil
primaryConn := c.primaryConn
c.connMutex.RUnlock()
// If we now have a primary connection, retry the batch
if primaryAvailable && primaryConn != nil {
resp, sendErr = primaryConn.Send(timeoutCtx, transport.NewRequest(transport.TypeBatchWrite, reqData))
}
}
}
}
if sendErr != nil {
return false, fmt.Errorf("failed to send request: %w", sendErr)
}
var batchResp struct {
Success bool `json:"success"`
}
if err := json.Unmarshal(resp.Payload(), &batchResp); err != nil {
return false, fmt.Errorf("failed to unmarshal response: %w", err)
}
return batchResp.Success, nil
}
// GetStats retrieves database statistics
func (c *Client) GetStats(ctx context.Context) (*Stats, error) {
if !c.IsConnected() {
return nil, errors.New("not connected to server")
}
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
// GetStats doesn't require a payload
resp, err := c.client.Send(timeoutCtx, transport.NewRequest(transport.TypeGetStats, nil))
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
var statsResp struct {
KeyCount int64 `json:"key_count"`
StorageSize int64 `json:"storage_size"`
MemtableCount int32 `json:"memtable_count"`
SstableCount int32 `json:"sstable_count"`
WriteAmplification float64 `json:"write_amplification"`
ReadAmplification float64 `json:"read_amplification"`
}
if err := json.Unmarshal(resp.Payload(), &statsResp); err != nil {
return nil, fmt.Errorf("failed to unmarshal response: %w", err)
}
return &Stats{
KeyCount: statsResp.KeyCount,
StorageSize: statsResp.StorageSize,
MemtableCount: statsResp.MemtableCount,
SstableCount: statsResp.SstableCount,
WriteAmplification: statsResp.WriteAmplification,
ReadAmplification: statsResp.ReadAmplification,
}, nil
}
// Compact triggers compaction of the database
func (c *Client) Compact(ctx context.Context, force bool) (bool, error) {
if !c.IsConnected() {
return false, errors.New("not connected to server")
}
req := struct {
Force bool `json:"force"`
}{
Force: force,
}
reqData, err := json.Marshal(req)
if err != nil {
return false, fmt.Errorf("failed to marshal request: %w", err)
}
timeoutCtx, cancel := context.WithTimeout(ctx, c.options.RequestTimeout)
defer cancel()
resp, err := c.client.Send(timeoutCtx, transport.NewRequest(transport.TypeCompact, reqData))
if err != nil {
return false, fmt.Errorf("failed to send request: %w", err)
}
var compactResp struct {
Success bool `json:"success"`
}
if err := json.Unmarshal(resp.Payload(), &compactResp); err != nil {
return false, fmt.Errorf("failed to unmarshal response: %w", err)
}
return compactResp.Success, nil
}
// Stats contains database statistics
type Stats struct {
KeyCount int64
StorageSize int64
MemtableCount int32
SstableCount int32
WriteAmplification float64
ReadAmplification float64
}
// GetNodeInfo returns information about the current node and replication topology
func (c *Client) GetReplicationInfo() (*NodeInfo, error) {
c.connMutex.RLock()
defer c.connMutex.RUnlock()
if c.nodeInfo == nil {
return nil, errors.New("replication information not available")
}
// Return a copy to avoid concurrent access issues
return &NodeInfo{
Role: c.nodeInfo.Role,
PrimaryAddr: c.nodeInfo.PrimaryAddr,
Replicas: c.nodeInfo.Replicas,
LastSequence: c.nodeInfo.LastSequence,
ReadOnly: c.nodeInfo.ReadOnly,
}, nil
}
// RefreshTopology updates the replication topology information
func (c *Client) RefreshTopology(ctx context.Context) error {
return c.discoverTopology(ctx)
}
// IsPrimary returns true if the connected node is a primary
func (c *Client) IsPrimary() bool {
c.connMutex.RLock()
defer c.connMutex.RUnlock()
return c.nodeInfo != nil && c.nodeInfo.Role == "primary"
}
// IsReplica returns true if the connected node is a replica
func (c *Client) IsReplica() bool {
c.connMutex.RLock()
defer c.connMutex.RUnlock()
return c.nodeInfo != nil && c.nodeInfo.Role == "replica"
}
// IsStandalone returns true if the connected node is standalone (not part of replication)
func (c *Client) IsStandalone() bool {
c.connMutex.RLock()
defer c.connMutex.RUnlock()
return c.nodeInfo == nil || c.nodeInfo.Role == "standalone"
}