textproto.CanonincalMIMEHeaderKey with memoization and GC

textproto.go

package textproto

import (
	"net/textproto"
	"runtime"
	"sync"
)

// CanonincalMIMEHeaderKey is like textproto.CanonicalMIMEHeaderKey but it
// memoizes results to avoid repeated allocations of the same string.
//
// The memoization is thread-safe and lazily cleaned up when the GC runs.
func CanonincalMIMEHeaderKey(k string) string {
	m.RLock()
	ck, ok := c[k]
	if ok {
		m.RUnlock()
		return ck
	}
	ck, ok = o[k]
	m.RUnlock()

	if !ok {
		ck = textproto.CanonicalMIMEHeaderKey(k)
	}

	m.Lock()
	if c == nil {
		c = make(map[string]string, len(o))
		if o == nil {
			defer runtime.SetFinalizer(new(canary), finalizer)
		}
	}
	c[k] = ck
	m.Unlock()

	return ck
}

var (
	m sync.RWMutex      // protects the following variables
	c map[string]string // current map
	o map[string]string // previous map
)

type canary [16]byte

func finalizer(x *canary) {
	go func() {
		m.Lock()
		if c != nil {
			o, c = c, nil
			defer runtime.SetFinalizer(x, finalizer) // rearm finalizer
		} else {
			o = nil
		}
		m.Unlock()
	}()
}