fixes from PR review

Author: asmyasnikov

retry/budget/budget.go

@@ -77,15 +77,15 @@ func (q *budget) Stop() {
 
 // Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
 func (q *budget) Acquire(ctx context.Context) error {
+	if err := ctx.Err(); err != nil {
+		return xerrors.WithStackTrace(err)
+	}
 	select {
+	case <-q.done:
+		return xerrors.WithStackTrace(errClosedBudget)
+	case <-q.quota:
+		return nil
 	case <-ctx.Done():
-		return ctx.Err()
-	default:
-		select {
-		case <-q.quota:
-			return nil
-		case <-ctx.Done():
-			return xerrors.WithStackTrace(ctx.Err())
-		}
+		return xerrors.WithStackTrace(ctx.Err())
 	}
 }

retry/budget/errors.go

@@ -6,7 +6,11 @@ import (
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
 )
 
-// ErrNoQuota is a special error for no quota provided by external retry budget
-//
-// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
-var ErrNoQuota = xerrors.Wrap(errors.New("no retry quota"))
+var (
+	// ErrNoQuota is a special error for no quota provided by external retry budget
+	//
+	// Experimental: https://github.com/ydb-platform/ydb-go-sdk/blob/master/VERSIONING.md#experimental
+	ErrNoQuota = xerrors.Wrap(errors.New("no retry quota"))
+
+	errClosedBudget = xerrors.Wrap(errors.New("retry budget closed"))
+)