<commit_before>package gkvlite

import (
	"encoding/binary"
	"fmt"
	"sync"
	"sync/atomic"
)

// A persistable node.
type node struct {
	numNodes, numBytes uint64
	item               itemLoc
	left, right        nodeLoc
	next               *node // For free-list tracking.
}

// A persistable node and its persistence location.
type nodeLoc struct {
	m sync.Mutex

	loc  *ploc    // *ploc - can be nil if node is dirty (not yet persisted).
	node *node    // *node - can be nil if node is not fetched into memory yet.
	next *nodeLoc // For free-list tracking.
}

var empty_nodeLoc = &nodeLoc{} // Sentinel.

func (nloc *nodeLoc) Loc() *ploc {
	nloc.m.Lock()
	defer nloc.m.Unlock()
	return nloc.loc
}

func (nloc *nodeLoc) setLoc(n *ploc) {
	nloc.m.Lock()
	defer nloc.m.Unlock()
	nloc.loc = n
}

func (nloc *nodeLoc) Node() *node {
	nloc.m.Lock()
	defer nloc.m.Unlock()
	return nloc.node
}

func (nloc *nodeLoc) setNode(n *node) {
	nloc.m.Lock()
	defer nloc.m.Unlock()
	nloc.node = n
}

func (nloc *nodeLoc) Copy(src *nodeLoc) *nodeLoc {
	if src == nil {
		return nloc.Copy(empty_nodeLoc)
	}
	newloc := src.Loc()
	newnode := src.Node()

	nloc.m.Lock()
	defer nloc.m.Unlock()
	nloc.loc = newloc
	nloc.node = newnode
	return nloc
}

func (nloc *nodeLoc) isEmpty() bool {
	return nloc == nil || (nloc.Loc().isEmpty() && nloc.Node() == nil)
}

func (nloc *nodeLoc) write(o *Store) error {
	if nloc != nil && nloc.Loc().isEmpty() {
		node := nloc.Node()
		if node == nil {
			return nil
		}
		offset := atomic.LoadInt64(&o.size)
		length := ploc_length + ploc_length + ploc_length + 8 + 8
		b := make([]byte, length)
		pos := 0
		pos = node.item.Loc().write(b, pos)
		pos = node.left.Loc().write(b, pos)
		pos = node.right.Loc().write(b, pos)
		binary.BigEndian.PutUint64(b[pos:pos+8], node.numNodes)
		pos += 8
		binary.BigEndian.PutUint64(b[pos:pos+8], node.numBytes)
		pos += 8
		if pos != length {
			return fmt.Errorf("nodeLoc.write() pos: %v didn't match length: %v",
				pos, length)
		}
		if _, err := o.file.WriteAt(b, offset); err != nil {
			return err
		}
		atomic.StoreInt64(&o.size, offset+int64(length))
		nloc.setLoc(&ploc{Offset: offset, Length: uint32(length)})
	}
	return nil
}

func (nloc *nodeLoc) read(o *Store) (n *node, err error) {
	if nloc == nil {
		return nil, nil
	}
	n = nloc.Node()
	if n != nil {
		return n, nil
	}
	loc := nloc.Loc()
	if loc.isEmpty() {
		return nil, nil
	}
	if loc.Length != uint32(ploc_length+ploc_length+ploc_length+8+8) {
		return nil, fmt.Errorf("unexpected node loc.Length: %v != %v",
			loc.Length, ploc_length+ploc_length+ploc_length+8+8)
	}
	b := make([]byte, loc.Length)
	if _, err := o.file.ReadAt(b, loc.Offset); err != nil {
		return nil, err
	}<commit_msg>Fix race in gkvlite introduced when removing `unsafe`<commit_after>package gkvlite

import (
	"encoding/binary"
	"fmt"
	"sync"
	"sync/atomic"
)

// A persistable node.
type node struct {
	numNodes, numBytes uint64
	item               itemLoc
	left, right        nodeLoc
	next               *node // For free-list tracking.
}

var nodeLocGL = sync.RWMutex{}

// A persistable node and its persistence location.
type nodeLoc struct {
	loc  *ploc    // *ploc - can be nil if node is dirty (not yet persisted).
	node *node    // *node - can be nil if node is not fetched into memory yet.
	next *nodeLoc // For free-list tracking.
}

var empty_nodeLoc = &nodeLoc{} // Sentinel.

func (nloc *nodeLoc) Loc() *ploc {
	nodeLocGL.RLock()
	defer nodeLocGL.RUnlock()
	return nloc.loc
}

func (nloc *nodeLoc) setLoc(n *ploc) {
	nodeLocGL.Lock()
	defer nodeLocGL.Unlock()
	nloc.loc = n
}

func (nloc *nodeLoc) Node() *node {
	nodeLocGL.RLock()
	defer nodeLocGL.RUnlock()
	return nloc.node
}

func (nloc *nodeLoc) setNode(n *node) {
	nodeLocGL.Lock()
	defer nodeLocGL.Unlock()
	nloc.node = n
}

func (nloc *nodeLoc) LocNode() (*ploc, *node) {
	nodeLocGL.RLock()
	defer nodeLocGL.RUnlock()
	return nloc.loc, nloc.node
}

func (nloc *nodeLoc) Copy(src *nodeLoc) *nodeLoc {
	if src == nil {
		return nloc.Copy(empty_nodeLoc)
	}

	nodeLocGL.Lock()
	defer nodeLocGL.Unlock()
	// NOTE: This trick only works because of the global lock. No reason to lock
	// src independently of nlock.
	nloc.loc = src.loc
	nloc.node = src.node
	return nloc
}

func (nloc *nodeLoc) isEmpty() bool {
	nodeLocGL.RLock()
	defer nodeLocGL.RUnlock()
	return nloc == nil || (nloc.loc.isEmpty() && nloc.node == nil)
}

func (nloc *nodeLoc) write(o *Store) error {
	loc, node := nloc.LocNode()

	if nloc != nil && loc.isEmpty() {
		if node == nil {
			return nil
		}
		offset := atomic.LoadInt64(&o.size)
		length := ploc_length + ploc_length + ploc_length + 8 + 8
		b := make([]byte, length)
		pos := 0
		pos = node.item.Loc().write(b, pos)
		pos = node.left.Loc().write(b, pos)
		pos = node.right.Loc().write(b, pos)
		binary.BigEndian.PutUint64(b[pos:pos+8], node.numNodes)
		pos += 8
		binary.BigEndian.PutUint64(b[pos:pos+8], node.numBytes)
		pos += 8
		if pos != length {
			return fmt.Errorf("nodeLoc.write() pos: %v didn't match length: %v",
				pos, length)
		}
		if _, err := o.file.WriteAt(b, offset); err != nil {
			return err
		}
		atomic.StoreInt64(&o.size, offset+int64(length))
		nloc.setLoc(&ploc{Offset: offset, Length: uint32(length)})
	}
	return nil
}

func (nloc *nodeLoc) read(o *Store) (n *node, err error) {
	if nloc == nil {
		return nil, nil
	}
	loc, n := nloc.LocNode()
	if n != nil {
		return n, nil
	}
	if loc.isEmpty() {
		return nil, nil
	}
	if loc.Length != uint32(ploc_length+ploc_length+ploc_length+8+8) {
		return nil, fmt.Errorf("unexpected node loc.Length: %v != %v",
			loc.Length, ploc_length+ploc_length+ploc_length+8+8)
	}
	b := make([]byte, loc.Length)
	if _, err := o.file.ReadAt(b, loc.Offset); err != nil {
		return nil, err
	}
{'fraction_non_alphanumeric': 0.15408750217883913, 'fraction_numerical': 0.009412584974725467, 'mean_word_length': 3.3142857142857145, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3758398', 'n_tokens_mistral': 2366, 'n_tokens_neox': 2249, 'n_words': 595}