Ref counter based on Go GC

counter.go

package main

import (
	"fmt"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

type token []byte

type item struct {
	n     int32
	empty chan struct{}
}

type counter struct {
	mu    sync.Mutex
	items map[interface{}]*item
	size  int
}

var closed = (func() <-chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
})()

func newCounter(tokenSize int) *counter {
	return &counter{
		items: make(map[interface{}]*item),
		size:  tokenSize,
	}
}

func (c *counter) refer(x interface{}) *token {
	c.mu.Lock()
	defer c.mu.Unlock()

	el := c.items[x]
	if el == nil {
		el = new(item)
		c.items[x] = el
	}
	if atomic.AddInt32(&el.n, 1) == 1 {
		el.empty = make(chan struct{})
	}

	t := token(make([]byte, c.size))
	runtime.SetFinalizer(&t, func(*token) {
		fmt.Println("fin")
		if atomic.AddInt32(&el.n, -1) != 0 {
			return
		}

		c.mu.Lock()
		defer c.mu.Unlock()

		if atomic.LoadInt32(&el.n) != 0 {
			return
		}

		close(el.empty)
		delete(c.items, x)
	})

	return &t
}

func (c *counter) collected(x interface{}) <-chan struct{} {
	c.mu.Lock()
	defer c.mu.Unlock()

	el := c.items[x]
	if el == nil {
		return closed
	}

	return el.empty
}

func main() {
	cnt := newCounter(1 << 12)

	var ts []*token
	ts = append(ts, cnt.refer(42))
	ts = append(ts, cnt.refer(42))
	ts = append(ts, cnt.refer(42))
	ts = append(ts, cnt.refer(42))

	go wait(cnt.collected(42))

	for i := range ts {
		ts[i] = nil
		time.Sleep(time.Millisecond * 100)
	}

	runtime.GC()

	cnt.refer(42)
	go wait(cnt.collected(42))

	runtime.GC()
}

func wait(done <-chan struct{}) {
	<-done
	fmt.Println("OK!")
}