84 lines
1.9 KiB
Go
84 lines
1.9 KiB
Go
|
package connection
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"fmt"
|
||
|
|
"log"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
pb "nexusrmm.local/agent/pkg/proto"
|
||
|
|
"google.golang.org/grpc"
|
||
|
|
"google.golang.org/grpc/connectivity"
|
||
|
|
"google.golang.org/grpc/credentials/insecure"
|
||
|
|
"google.golang.org/grpc/keepalive"
|
||
|
|
)
|
||
|
|
|
||
|
|
type GrpcClient struct {
|
||
|
|
conn *grpc.ClientConn
|
||
|
|
Client pb.AgentServiceClient
|
||
|
|
address string
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewGrpcClient(address string) (*GrpcClient, error) {
|
||
|
|
opts := []grpc.DialOption{
|
||
|
|
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
||
|
|
grpc.WithKeepaliveParams(keepalive.ClientParameters{
|
||
|
|
Time: 30 * time.Second,
|
||
|
|
Timeout: 10 * time.Second,
|
||
|
|
PermitWithoutStream: true,
|
||
|
|
}),
|
||
|
|
}
|
||
|
|
|
||
|
|
conn, err := grpc.NewClient(address, opts...)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
|
||
|
|
return &GrpcClient{
|
||
|
|
conn: conn,
|
||
|
|
Client: pb.NewAgentServiceClient(conn),
|
||
|
|
address: address,
|
||
|
|
}, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (g *GrpcClient) Close() {
|
||
|
|
if g.conn != nil {
|
||
|
|
g.conn.Close()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func ConnectWithRetry(address string, maxRetries int) (*GrpcClient, error) {
|
||
|
|
for i := 0; i < maxRetries; i++ {
|
||
|
|
client, err := NewGrpcClient(address)
|
||
|
|
if err != nil {
|
||
|
|
log.Printf("Client creation %d/%d failed: %v", i+1, maxRetries, err)
|
||
|
|
time.Sleep(backoffDuration(i))
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
|
|
client.conn.Connect()
|
||
|
|
state := client.conn.GetState()
|
||
|
|
client.conn.WaitForStateChange(ctx, state)
|
||
|
|
cancel()
|
||
|
|
|
||
|
|
newState := client.conn.GetState()
|
||
|
|
if newState == connectivity.Ready || newState == connectivity.Idle {
|
||
|
|
return client, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
log.Printf("Connection attempt %d/%d: state=%v", i+1, maxRetries, newState)
|
||
|
|
client.Close()
|
||
|
|
time.Sleep(backoffDuration(i))
|
||
|
|
}
|
||
|
|
return nil, fmt.Errorf("failed to connect after %d attempts", maxRetries)
|
||
|
|
}
|
||
|
|
|
||
|
|
func backoffDuration(attempt int) time.Duration {
|
||
|
|
secs := 2 << attempt
|
||
|
|
if secs > 30 {
|
||
|
|
secs = 30
|
||
|
|
}
|
||
|
|
return time.Duration(secs) * time.Second
|
||
|
|
}
|