这是indexloc提供的服务，不要输入任何密码
Skip to content

feat: implement graceful shutdown for subscription client #171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
May 12, 2025
Merged

feat: implement graceful shutdown for subscription client #171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
10b3af0
feat: implement graceful shutdown for subscription client
rbroggi May 12, 2025
8d7f1c8
feat: add waitForConnectionState helper for subscription client tests
rbroggi May 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions subscription.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1165,6 +1165,13 @@ func (sc *SubscriptionClient) RunWithContext(ctx context.Context) error {
for {
select {
case <-ctx.Done():
// Check if the cancellation came from the parent context
if errors.Is(ctx.Err(), context.Canceled) {
// Parent context was canceled, close gracefully without error
_ = sc.close(subContext)
return nil
}
// Internal cancellation, return error
return sc.close(subContext)
case e := <-sc.errorChan:
if sc.getClientStatus() == scStatusClosing {
Expand Down
88 changes: 88 additions & 0 deletions subscription_graphql_ws_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -567,3 +567,91 @@ func TestSubscription_closeThenRun(t *testing.T) {
t.Fatalf("got error: %v, want: nil", err)
}
}

// waitForConnectionState waits for the subscription client to reach a specific connection state
func waitForConnectionState(t *testing.T, sc *SubscriptionClient, tickerDuration time.Duration, timeout time.Duration, checkFn func() bool) bool {
Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

utility method to avoid waiting for an established amount and delaying test execution.

t.Helper()
deadline := time.Now().Add(timeout)
ticker := time.NewTicker(tickerDuration)
defer ticker.Stop()

for {
select {
case <-ticker.C:
if checkFn() {
return true
}
if time.Now().After(deadline) {
return false
}
}
}
}

func TestRunWithContext_GracefulShutdown(t *testing.T) {
subscriptionClient := NewSubscriptionClient(fmt.Sprintf("%s/v1/graphql", hasuraTestHost)).
WithConnectionParams(map[string]interface{}{
"headers": map[string]string{
"x-hasura-admin-secret": hasuraTestAdminSecret,
},
}).
WithProtocol(GraphQLWS).
WithLog(log.Println)

ctx, cancel := context.WithCancel(context.Background())

// Simulate a subscription
var sub struct {
Users []struct {
ID int `graphql:"id"`
Name string `graphql:"name"`
} `graphql:"user(order_by: { id: desc }, limit: 5)"`
}

_, err := subscriptionClient.Subscribe(sub, nil, func(data []byte, e error) error {
if e != nil {
t.Fatalf("got error: %v, want: nil", e)
return nil
}
return nil
})
if err != nil {
t.Fatalf("got error: %v, want: nil", err)
}

// Run the subscription client in a separate goroutine and collect error
done := make(chan struct{})
var runErr error
go func() {
runErr = subscriptionClient.RunWithContext(ctx)
close(done)
}()

// Wait for the client to establish connection
if !waitForConnectionState(t, subscriptionClient, 100*time.Millisecond, 5*time.Second, func() bool {
session := subscriptionClient.getCurrentSession()
return session != nil && session.GetAcknowledge()
}) {
t.Fatal("timeout waiting for connection to be established")
}

// Cancel the parent context to trigger graceful shutdown
cancel()

// Wait for the client to shut down
if !waitForConnectionState(t, subscriptionClient, 100*time.Millisecond, 5*time.Second, func() bool {
return subscriptionClient.getCurrentSession() == nil
}) {
t.Fatal("timeout waiting for connection to be closed")
}

// Wait for the Run goroutine to finish
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("timeout waiting for Run goroutine to finish")
}
if runErr != nil {
t.Errorf("got error: %v, want: nil", err)
}
}
Loading