@@ -137,7 +137,7 @@ func RequireNodeSecurityMember(c *cli.Command) error {
137137// Service synchronization
138138//
139139
140- func WaitNodeHdPassword (c * cli.Command , verbose bool ) error {
140+ func WaitNodeHdPassword (ctx context. Context , c * cli.Command , verbose bool ) error {
141141 for {
142142 nodePasswordSet , err := getNodeHdPasswordSet (c )
143143 if err != nil {
@@ -149,12 +149,14 @@ func WaitNodeHdPassword(c *cli.Command, verbose bool) error {
149149 if verbose {
150150 log .Printf ("The node password has not been set, retrying in %s...\n " , checkNodePasswordInterval .String ())
151151 }
152- time .Sleep (checkNodePasswordInterval )
152+ if err := sleepCtx (ctx , checkNodePasswordInterval ); err != nil {
153+ return err
154+ }
153155 }
154156}
155157
156- func WaitNodeHdWallet (c * cli.Command , verbose bool ) error {
157- if err := WaitNodeHdPassword (c , verbose ); err != nil {
158+ func WaitNodeHdWallet (ctx context. Context , c * cli.Command , verbose bool ) error {
159+ if err := WaitNodeHdPassword (ctx , c , verbose ); err != nil {
158160 return err
159161 }
160162 for {
@@ -168,7 +170,9 @@ func WaitNodeHdWallet(c *cli.Command, verbose bool) error {
168170 if verbose {
169171 log .Printf ("The node wallet has not been initialized, retrying in %s...\n " , checkNodeWalletInterval .String ())
170172 }
171- time .Sleep (checkNodeWalletInterval )
173+ if err := sleepCtx (ctx , checkNodeWalletInterval ); err != nil {
174+ return err
175+ }
172176 }
173177}
174178
@@ -182,7 +186,7 @@ func WaitBeaconClientSynced(c *cli.Command, verbose bool) error {
182186 return err
183187}
184188
185- func WaitRocketStorage (c * cli.Command , verbose bool ) error {
189+ func WaitRocketStorage (ctx context. Context , c * cli.Command , verbose bool ) error {
186190 if err := WaitEthClientSynced (c , verbose ); err != nil {
187191 return err
188192 }
@@ -197,16 +201,18 @@ func WaitRocketStorage(c *cli.Command, verbose bool) error {
197201 if verbose {
198202 log .Printf ("The Rocket Pool storage contract was not found, retrying in %s...\n " , checkRocketStorageInterval .String ())
199203 }
200- time .Sleep (checkRocketStorageInterval )
204+ if err := sleepCtx (ctx , checkRocketStorageInterval ); err != nil {
205+ return err
206+ }
201207 }
202208}
203209
204210// This check makes calls to GetHdWallet instead of GetWallet as it's used in node and watchtower
205- func WaitNodeRegistered (c * cli.Command , verbose bool ) error {
206- if err := WaitNodeHdWallet (c , verbose ); err != nil {
211+ func WaitNodeRegistered (ctx context. Context , c * cli.Command , verbose bool ) error {
212+ if err := WaitNodeHdWallet (ctx , c , verbose ); err != nil {
207213 return err
208214 }
209- if err := WaitRocketStorage (c , verbose ); err != nil {
215+ if err := WaitRocketStorage (ctx , c , verbose ); err != nil {
210216 return err
211217 }
212218 for {
@@ -220,7 +226,21 @@ func WaitNodeRegistered(c *cli.Command, verbose bool) error {
220226 if verbose {
221227 log .Printf ("The node is not registered with Rocket Pool, retrying in %s...\n " , checkNodeRegisteredInterval .String ())
222228 }
223- time .Sleep (checkNodeRegisteredInterval )
229+ if err := sleepCtx (ctx , checkNodeRegisteredInterval ); err != nil {
230+ return err
231+ }
232+ }
233+ }
234+
235+ // sleepCtx sleeps for d, returning ctx.Err() if the context is cancelled first.
236+ func sleepCtx (ctx context.Context , d time.Duration ) error {
237+ timer := time .NewTimer (d )
238+ defer timer .Stop ()
239+ select {
240+ case <- ctx .Done ():
241+ return ctx .Err ()
242+ case <- timer .C :
243+ return nil
224244 }
225245}
226246
0 commit comments