Custom Defer Implementation

custom_defer.go

package main

import (
	"errors"
	"fmt"
)

/*
Let's try to simulate defer manually.

We have the following situation:

Setup: func() {
    // setup op 1
    // setup op 2
}

Teardown: func() {
    // teardown op 2
    // teardown op 1
}

We want the teardown function to only run the teardown operations for
successful setup operations.

Example:

Setup: func() {
    // setup op 1: ok
    // setup op 2: error
}

Teardown: func() {
    // teardown op 2: *SKIP*
    // teardown op 1: RUN
}
*/

func main() {
	h := Happy{}
	h.Run()
	fmt.Println("db after Happy:", db)

	u := Unhappy{}
	u.Run()
	fmt.Println("db after Unhappy:", db)
}

// simulate a store of data
var db map[string]string = make(map[string]string)

type DeferredCleaner interface {
	Setup()
	Teardown()
	Run()
}

type Cleaner struct {
	CleanupStack []func()
}

func (c *Cleaner) Cleanup() {
	for i := len(c.CleanupStack) - 1; i >= 0; i-- {
		fn := c.CleanupStack[i]
		fn()
	}
}

func (c *Cleaner) Defer(fn func()) {
	c.CleanupStack = append(c.CleanupStack, fn)
}

type Happy struct {
}

var _ DeferredCleaner = (*Happy)(nil)

func (h *Happy) Setup() {
	okOperation("1", "one")
	okOperation("2", "two")
}

func (h *Happy) Teardown() {
	resetValue("2")
	resetValue("1")
}

func (h *Happy) Run() {
	h.Setup()
	h.Teardown()
}

type Unhappy struct {
	Cleaner
}

var _ DeferredCleaner = (*Unhappy)(nil)

func (u *Unhappy) Setup() {
	var err error
	three := "3"
	four := "4"
	five := "5"
	err = okOperation(three, "three")
	if err == nil {
		u.Defer(func() {
			resetValue(three)
		})
	}
	err = okOperation(four, "four")
	if err == nil {
		u.Defer(func() {
			resetValue(four)
		})
	}
	err = failOperation(five, "five")
	if err == nil {
		u.Defer(func() {
			resetValue(five)
		})
	}
}

func (u *Unhappy) Teardown() {
	u.Cleanup()
}

func (u *Unhappy) Run() {
	u.Setup()
	u.Teardown()
}

func okOperation(key string, val string) error {
	db[key] = val // act like we saved something
	return nil
}

func failOperation(key string, val string) error {
	return errors.New("failed operation")
}

func resetValue(key string) {
	fmt.Print("Delete key:", key)
	_, exists := db[key]
	if exists {
		fmt.Println(" ok")
		delete(db, key)
	} else {
		fmt.Println(" ROFL! it doesn't exist here!")
		panic("")
	}
}