cleanup.go 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859
  1. package framework
  2. import (
  3. "sync"
  4. )
  5. // CleanupActionHandle is an integer pointer type for handling cleanup action
  6. type CleanupActionHandle *int
  7. type cleanupFuncHandle struct {
  8. actionHandle CleanupActionHandle
  9. actionHook func()
  10. }
  11. var cleanupActionsLock sync.Mutex
  12. var cleanupHookList = []cleanupFuncHandle{}
  13. // AddCleanupAction installs a function that will be called in the event of the
  14. // whole test being terminated. This allows arbitrary pieces of the overall
  15. // test to hook into SynchronizedAfterSuite().
  16. // The hooks are called in last-in-first-out order.
  17. func AddCleanupAction(fn func()) CleanupActionHandle {
  18. p := CleanupActionHandle(new(int))
  19. cleanupActionsLock.Lock()
  20. defer cleanupActionsLock.Unlock()
  21. c := cleanupFuncHandle{actionHandle: p, actionHook: fn}
  22. cleanupHookList = append([]cleanupFuncHandle{c}, cleanupHookList...)
  23. return p
  24. }
  25. // RemoveCleanupAction removes a function that was installed by
  26. // AddCleanupAction.
  27. func RemoveCleanupAction(p CleanupActionHandle) {
  28. cleanupActionsLock.Lock()
  29. defer cleanupActionsLock.Unlock()
  30. for i, item := range cleanupHookList {
  31. if item.actionHandle == p {
  32. cleanupHookList = append(cleanupHookList[:i], cleanupHookList[i+1:]...)
  33. break
  34. }
  35. }
  36. }
  37. // RunCleanupActions runs all functions installed by AddCleanupAction. It does
  38. // not remove them (see RemoveCleanupAction) but it does run unlocked, so they
  39. // may remove themselves.
  40. func RunCleanupActions() {
  41. list := []func(){}
  42. func() {
  43. cleanupActionsLock.Lock()
  44. defer cleanupActionsLock.Unlock()
  45. for _, p := range cleanupHookList {
  46. list = append(list, p.actionHook)
  47. }
  48. }()
  49. // Run unlocked.
  50. for _, fn := range list {
  51. fn()
  52. }
  53. }