add shortcut method

Author: hsldymq

keeper.go

@@ -200,6 +200,11 @@ func (k *ShutdownKeeper) DeadlineContext() context.Context {
     return k.deadlineCtx
 }
 
+// DoOnShutdown is just a shortcut for HoldToken's DoOnShutdown method.
+func (k *ShutdownKeeper) DoOnShutdown(f func(ctx context.Context)) {
+    k.AllocHoldToken().DoOnShutdown(f)
+}
+
 func (k *ShutdownKeeper) doAllocToken(allocFunc func(releaseCallback func()) HoldToken) HoldToken {
     atomic.AddInt32(&k.holdTokenNum, 1)
     return allocFunc(sync.OnceFunc(k.doReleaseToken))

keeper_test.go

@@ -65,12 +65,12 @@ func TestShutdownKeeper_HoldTokens(t *testing.T) {
 
     start := time.Now()
     var actual int32
-    keeper.AllocHoldToken().DoOnShutdown(func(_ context.Context) {
+    keeper.DoOnShutdown(func(_ context.Context) {
         time.Sleep(time.Second)
         atomic.AddInt32(&actual, 1)
     })
 
-    keeper.AllocHoldToken().DoOnShutdown(func(_ context.Context) {
+    keeper.DoOnShutdown(func(_ context.Context) {
         time.Sleep(500 * time.Millisecond)
         atomic.AddInt32(&actual, 1)
     })