Add context to client/server APIs and timeout configurations

nexcode · nexcode · commit 4935b971ca36 · 2026-03-30T21:33:13.000Z
diff --git a/README.md b/README.md
@@ -85,7 +85,7 @@ func main() {
 
 	proto.RegisterSumServer(server.Server(), &sumServer{})
 
-	if err = server.Serve(); err != nil {
+	if err = server.Serve(context.Background()); err != nil {
 		panic(err)
 	}
 }
@@ -128,7 +128,7 @@ func main() {
 		panic(err)
 	}
 
-	client, err := rpcp.NewClient("myServerName")
+	client, err := rpcp.NewClient(context.Background(), "myServerName")
 	if err != nil {
 		panic(err)
 	}
diff --git a/examples/attributes/client/main.go b/examples/attributes/client/main.go
@@ -61,7 +61,7 @@ func main() {
 		}
 	}()
 
-	client, err := rpcp.NewClient("myServerName", rpcplatform.ClientOptions.MaxActiveServers(2))
+	client, err := rpcp.NewClient(context.Background(), "myServerName", rpcplatform.ClientOptions.MaxActiveServers(2))
 	if err != nil {
 		panic(err)
 	}
diff --git a/examples/attributes/server/main.go b/examples/attributes/server/main.go
@@ -72,7 +72,7 @@ func main() {
 
 	proto.RegisterSumServer(server.Server(), &sumServer{})
 
-	if err = server.Serve(); err != nil {
+	if err = server.Serve(context.Background()); err != nil {
 		panic(err)
 	}
 }
diff --git a/examples/opentelemetry/client/main.go b/examples/opentelemetry/client/main.go
@@ -64,7 +64,7 @@ func main() {
 		panic(err)
 	}
 
-	client, err := rpcp.NewClient("myServerName")
+	client, err := rpcp.NewClient(context.Background(), "myServerName")
 	if err != nil {
 		panic(err)
 	}
diff --git a/examples/opentelemetry/server/main.go b/examples/opentelemetry/server/main.go
@@ -85,7 +85,7 @@ func main() {
 
 	proto.RegisterSumServer(server.Server(), &sumServer{})
 
-	if err = server.Serve(); err != nil {
+	if err = server.Serve(context.Background()); err != nil {
 		panic(err)
 	}
 }
diff --git a/examples/quickstart/client/main.go b/examples/quickstart/client/main.go
@@ -48,7 +48,7 @@ func main() {
 		panic(err)
 	}
 
-	client, err := rpcp.NewClient("myServerName")
+	client, err := rpcp.NewClient(context.Background(), "myServerName")
 	if err != nil {
 		panic(err)
 	}
diff --git a/examples/quickstart/server/main.go b/examples/quickstart/server/main.go
@@ -69,7 +69,7 @@ func main() {
 
 	proto.RegisterSumServer(server.Server(), &sumServer{})
 
-	if err = server.Serve(); err != nil {
+	if err = server.Serve(context.Background()); err != nil {
 		panic(err)
 	}
 }
diff --git a/internal/config/client.go b/internal/config/client.go
@@ -17,14 +17,19 @@
 package config
 
 import (
+	"time"
+
 	"google.golang.org/grpc"
 )
 
 func NewClient() *Client {
-	return &Client{}
+	return &Client{
+		EtcdClientTimeout: 5 * time.Second,
+	}
 }
 
 type Client struct {
-	MaxActiveServers int
-	GRPCOptions      []grpc.DialOption
+	MaxActiveServers  int
+	EtcdClientTimeout time.Duration
+	GRPCOptions       []grpc.DialOption
 }
diff --git a/internal/config/server.go b/internal/config/server.go
@@ -17,18 +17,25 @@
 package config
 
 import (
+	"time"
+
 	"github.com/nexcode/rpcplatform/internal/attributes"
 	"google.golang.org/grpc"
 )
 
 func NewServer() *Server {
 	return &Server{
-		Attributes: attributes.New(),
+		EtcdClientTimeout: 5 * time.Second,
+		EtcdLeaseTimeout:  5 * time.Second,
+		Attributes:        attributes.New(),
 	}
 }
 
 type Server struct {
-	PublicAddr  string
-	Attributes  *attributes.Attributes
-	GRPCOptions []grpc.ServerOption
+	PublicAddr        string
+	StopTimeout       time.Duration
+	EtcdClientTimeout time.Duration
+	EtcdLeaseTimeout  time.Duration
+	Attributes        *attributes.Attributes
+	GRPCOptions       []grpc.ServerOption
 }
diff --git a/internal/gears/сontextеimeout.go b/internal/gears/сontextеimeout.go
@@ -0,0 +1,14 @@
+package gears
+
+import (
+	"context"
+	"time"
+)
+
+func ContextTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
+	if timeout > 0 {
+		return context.WithTimeout(ctx, timeout)
+	}
+
+	return ctx, func() {}
+}
diff --git a/internal/options/client.go b/internal/options/client.go
@@ -17,6 +17,8 @@
 package options
 
 import (
+	"time"
+
 	"github.com/nexcode/rpcplatform/internal/config"
 	"google.golang.org/grpc"
 )
@@ -31,6 +33,14 @@ func (Client) MaxActiveServers(count int) func(*config.Client) {
 	}
 }
 
+// EtcdClientTimeout sets the timeout duration for server-side etcd client operations.
+// The default value is 5 seconds.
+func (Client) EtcdClientTimeout(timeout time.Duration) func(*config.Client) {
+	return func(c *config.Client) {
+		c.EtcdClientTimeout = timeout
+	}
+}
+
 // GRPCOptions adds gRPC dial options to the client.
 func (Client) GRPCOptions(options ...grpc.DialOption) func(*config.Client) {
 	return func(c *config.Client) {
diff --git a/internal/options/server.go b/internal/options/server.go
@@ -17,6 +17,8 @@
 package options
 
 import (
+	"time"
+
 	"github.com/nexcode/rpcplatform/internal/attributes"
 	"github.com/nexcode/rpcplatform/internal/config"
 	"google.golang.org/grpc"
@@ -31,6 +33,31 @@ func (Server) PublicAddr(publicAddr string) func(*config.Server) {
 	}
 }
 
+// StopTimeout sets the maximum duration the server will wait for a graceful stop.
+// If the timeout is greater than zero, a graceful stop is applied; otherwise, it stops immediately.
+// The default value is 0.
+func (Server) StopTimeout(timeout time.Duration) func(*config.Server) {
+	return func(c *config.Server) {
+		c.StopTimeout = timeout
+	}
+}
+
+// EtcdClientTimeout sets the timeout duration for server-side etcd client operations.
+// The default value is 5 seconds.
+func (Server) EtcdClientTimeout(timeout time.Duration) func(*config.Server) {
+	return func(c *config.Server) {
+		c.EtcdClientTimeout = timeout
+	}
+}
+
+// EtcdLeaseTimeout sets the Time-To-Live (TTL) duration for the etcd lease.
+// The default value is 5 seconds.
+func (Server) EtcdLeaseTimeout(timeout time.Duration) func(*config.Server) {
+	return func(c *config.Server) {
+		c.EtcdLeaseTimeout = timeout
+	}
+}
+
 // Attributes sets server attributes that are applied by the server and accessible via the Lookup method.
 func (Server) Attributes(attributes *attributes.Attributes) func(*config.Server) {
 	return func(c *config.Server) {
diff --git a/rpcplatform_newclient.go b/rpcplatform_newclient.go
@@ -27,10 +27,11 @@ import (
 	"github.com/nexcode/rpcplatform/internal/gears"
 	"github.com/nexcode/rpcplatform/internal/resolver"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/connectivity"
 )
 
 // NewClient creates a new client connecting to the specified server name.
-func (p *RPCPlatform) NewClient(target string, options ...ClientOption) (*Client, error) {
+func (p *RPCPlatform) NewClient(ctx context.Context, target string, options ...ClientOption) (*Client, error) {
 	if target == "" || strings.Contains(target, "/") {
 		return nil, fmt.Errorf("%q: target is empty or contains «/»: %w", target, ErrInvalidTargetName)
 	}
@@ -52,24 +53,22 @@ func (p *RPCPlatform) NewClient(target string, options ...ClientOption) (*Client
 		config:   config,
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
+	timer := time.AfterFunc(config.EtcdClientTimeout, func() { cancel() })
 
 	serverInfoTree, err := p.Lookup(ctx, target, true)
-	if err != nil {
-		return nil, err
-	}
 
-	select {
-	case serverInfoTree := <-serverInfoTree:
-		c.updateState(true, serverInfoTree)
-	case <-time.After(4 * time.Second):
-		cancel()
+	if !timer.Stop() {
+		<-ctx.Done()
+		return nil, ctx.Err()
 	}
 
-	if err = ctx.Err(); err != nil {
+	if err != nil {
 		return nil, err
 	}
 
+	c.updateState(true, <-serverInfoTree)
+
 	config.GRPCOptions = append(config.GRPCOptions,
 		grpc.WithResolvers(c.resolver),
 		grpc.WithDefaultServiceConfig(`{"loadBalancingConfig":[{"`+balancer.Name+`":{}}]}`),
@@ -90,10 +89,27 @@ func (p *RPCPlatform) NewClient(target string, options ...ClientOption) (*Client
 	}
 
 	go func() {
+		defer c.client.Close()
+
 		for serverInfoTree := range serverInfoTree {
 			c.updateState(false, serverInfoTree)
 		}
 	}()
 
+	go func() {
+		state := c.client.GetState()
+
+		for {
+			if !c.client.WaitForStateChange(ctx, state) {
+				return
+			}
+
+			if state = c.client.GetState(); state == connectivity.Shutdown {
+				cancel()
+				return
+			}
+		}
+	}()
+
 	return c, nil
 }
diff --git a/rpcplatform_test.go b/rpcplatform_test.go
@@ -183,7 +183,7 @@ func TestRPCPlatform_Lookup(t *testing.T) {
 	defer server.Server().Stop()
 
 	go func() {
-		if err := server.Serve(); err != nil {
+		if err := server.Serve(context.Background()); err != nil {
 			t.Errorf("Serve() failed: %v", err)
 		}
 	}()
@@ -305,7 +305,7 @@ func TestRPCPlatform_NewClient(t *testing.T) {
 				t.Fatalf("New() failed: %v", err)
 			}
 
-			client, err := rpcp.NewClient(tt.input.target, tt.input.clientOptions...)
+			client, err := rpcp.NewClient(context.Background(), tt.input.target, tt.input.clientOptions...)
 			if err != nil {
 				t.Fatalf("NewClient() failed: %v", err)
 			}
@@ -503,7 +503,7 @@ func TestRPCPlatform_NewServer(t *testing.T) {
 				server.Server().Stop()
 			}()
 
-			if err := server.Serve(); err != nil {
+			if err := server.Serve(context.Background()); err != nil {
 				t.Errorf("Serve() failed: %v", err)
 			}
 		})
diff --git a/server_serve.go b/server_serve.go
@@ -18,36 +18,45 @@ package rpcplatform
 
 import (
 	"context"
-	"fmt"
+	"log"
 	"time"
 
 	"github.com/nexcode/rpcplatform/internal/attributes"
+	"github.com/nexcode/rpcplatform/internal/gears"
 	etcd "go.etcd.io/etcd/client/v3"
 )
 
 // Serve starts the gRPC server and blocks until it exits or an error occurs.
-func (s *Server) Serve() error {
-	global, cancel := context.WithCancel(context.Background())
+func (s *Server) Serve(ctx context.Context) error {
+	path := s.name + "/" + s.id
+	attributes := attributes.Values(s.config.Attributes)
 
-	defer func() {
-		cancel()
-	}()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
 	go func() {
-		path := s.name + "/" + s.id
-		attributes := attributes.Values(s.config.Attributes)
+		defer func() {
+			timer := time.AfterFunc(s.config.StopTimeout, func() {
+				s.Server().Stop()
+			})
+
+			if s.config.StopTimeout > 0 {
+				s.Server().GracefulStop()
+				timer.Stop()
+			}
+		}()
 
 		for {
-			if global.Err() != nil {
-				break
+			if ctx.Err() != nil {
+				return
 			}
 
-			ctx, cancel := context.WithTimeout(global, 4*time.Second)
-			lease, err := s.etcd.Grant(ctx, 4)
-			cancel()
+			ctxTimeout, cancelTimeout := gears.ContextTimeout(ctx, s.config.EtcdClientTimeout)
+			lease, err := s.etcd.Grant(ctxTimeout, int64(s.config.EtcdLeaseTimeout.Seconds()))
+			cancelTimeout()
 
 			if err != nil {
-				fmt.Println(err)
+				log.Println(err)
 				continue
 			}
 
@@ -63,22 +72,22 @@ func (s *Server) Serve() error {
 				ops = append(ops, etcd.OpPut(path+"/"+attributes[i], attributes[i+1], etcd.WithLease(lease.ID)))
 			}
 
-			ctx, cancel = context.WithTimeout(global, 4*time.Second)
-			resp, err := s.etcd.Txn(ctx).Then(ops...).Commit()
-			cancel()
+			ctxTimeout, cancelTimeout = gears.ContextTimeout(ctx, s.config.EtcdClientTimeout)
+			resp, err := s.etcd.Txn(ctxTimeout).Then(ops...).Commit()
+			cancelTimeout()
 
 			if err != nil {
-				fmt.Println(err)
+				log.Println(err)
 				continue
 			}
 
 			if !resp.Succeeded {
 				continue
 			}
 
-			keepAlive, err := s.etcd.KeepAlive(global, lease.ID)
+			keepAlive, err := s.etcd.KeepAlive(ctx, lease.ID)
 			if err != nil {
-				fmt.Println(err)
+				log.Println(err)
 				continue
 			}
 

Original file line number	Diff line number	Diff line change
`@@ -85,7 +85,7 @@ func main() {`
`85`	`85`
`86`	`86`	`proto.RegisterSumServer(server.Server(), &sumServer{})`
`87`	`87`
`88`		`- if err = server.Serve(); err != nil {`
	`88`	`+ if err = server.Serve(context.Background()); err != nil {`
`89`	`89`	`panic(err)`
`90`	`90`	`}`
`91`	`91`	`}`
`@@ -128,7 +128,7 @@ func main() {`
`128`	`128`	`panic(err)`
`129`	`129`	`}`
`130`	`130`
`131`		`- client, err := rpcp.NewClient("myServerName")`
	`131`	`+ client, err := rpcp.NewClient(context.Background(), "myServerName")`
`132`	`132`	`if err != nil {`
`133`	`133`	`panic(err)`
`134`	`134`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ func main() {`
`61`	`61`	`}`
`62`	`62`	`}()`
`63`	`63`
`64`		`- client, err := rpcp.NewClient("myServerName", rpcplatform.ClientOptions.MaxActiveServers(2))`
	`64`	`+ client, err := rpcp.NewClient(context.Background(), "myServerName", rpcplatform.ClientOptions.MaxActiveServers(2))`
`65`	`65`	`if err != nil {`
`66`	`66`	`panic(err)`
`67`	`67`	`}`
Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ func main() {`
`72`	`72`
`73`	`73`	`proto.RegisterSumServer(server.Server(), &sumServer{})`
`74`	`74`
`75`		`- if err = server.Serve(); err != nil {`
	`75`	`+ if err = server.Serve(context.Background()); err != nil {`
`76`	`76`	`panic(err)`
`77`	`77`	`}`
`78`	`78`	`}`
Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ func main() {`
`64`	`64`	`panic(err)`
`65`	`65`	`}`
`66`	`66`
`67`		`- client, err := rpcp.NewClient("myServerName")`
	`67`	`+ client, err := rpcp.NewClient(context.Background(), "myServerName")`
`68`	`68`	`if err != nil {`
`69`	`69`	`panic(err)`
`70`	`70`	`}`
Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ func main() {`
`48`	`48`	`panic(err)`
`49`	`49`	`}`
`50`	`50`
`51`		`- client, err := rpcp.NewClient("myServerName")`
	`51`	`+ client, err := rpcp.NewClient(context.Background(), "myServerName")`
`52`	`52`	`if err != nil {`
`53`	`53`	`panic(err)`
`54`	`54`	`}`
Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ func main() {`
`69`	`69`
`70`	`70`	`proto.RegisterSumServer(server.Server(), &sumServer{})`
`71`	`71`
`72`		`- if err = server.Serve(); err != nil {`
	`72`	`+ if err = server.Serve(context.Background()); err != nil {`
`73`	`73`	`panic(err)`
`74`	`74`	`}`
`75`	`75`	`}`