<commit_before>// Copyright © 2016 Zenly <<EMAIL>>.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package protein

import (
	"reflect"
	"sync"

	"go.uber.org/zap"

	"github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
	"github.com/pkg/errors"
	"github.com/znly/protein/failure"
	"github.com/znly/protein/protoscan"
)

// -----------------------------------------------------------------------------

// SchemaMap is a thread-safe mapping & reverse-mapping of `ProtobufSchema`s.
//
// It atomically maintains two data-structures in parallel:
// a map of schemaUIDs to `ProtobufSchema`s, and
// a map of fully-qualified schema names to schemaUIDs.
//
// The `SchemaMap` is the main data-structure behind a `Transcoder`, used to
// store and retrieve every `ProtobufSchema`s that have been cached locally.
type SchemaMap struct {
	lock *sync.RWMutex
	// map[schemaUID]ProtobufSchema
	schemaMap map[string]*ProtobufSchema
	// reverse-mapping of fully-qualified names to UIDs
	// map[FQNname][]{schemaUID}
	revmap map[string][]string
}

// NewSchemaMap returns a new SchemaMap with its maps & locks pre-allocated.
func NewSchemaMap() *SchemaMap {
	return &SchemaMap{
		lock:      &sync.RWMutex{},
		schemaMap: make(map[string]*ProtobufSchema, 0),
		revmap:    make(map[string][]string, 0),
	}
}

// Add walks over the given map of `schemas` and add them to the `SchemaMap`
// while making sure to atomically maintain both the internal map and reverse-map.
func (sm *SchemaMap) Add(schemas map[string]*ProtobufSchema) *SchemaMap {
	sm.lock.Lock()
	for _, s := range schemas {
		zap.L().Debug("schema found",
			zap.String("uid", s.SchemaUID), zap.String("fq-name", s.FQName),
		)
		sm.schemaMap[s.SchemaUID] = s
		sm.revmap[s.FQName] = append(sm.revmap[s.FQName], s.SchemaUID)
	}
	sm.lock.Unlock() // avoid defer()
	return sm
}

// Size returns the number of schemas stored in the `SchemaMap`.
func (sm *SchemaMap) Size() int {
	sm.lock.RLock()
	sz := len(sm.schemaMap)
	sm.lock.RUnlock() // avoid defer()
	return sz
}

// ForEach applies the specified function `f` to every entry in the `SchemaMap`.
//
// ForEach is guaranteed to see a consistent view of the internal mapping.
func (sm *SchemaMap) ForEach(f func(s *ProtobufSchema) error) error {
	sm.lock.RLock()
	for _, s := range sm.schemaMap {
		if err := f(s); err != nil {
			sm.lock.RUnlock() // avoid defer()
			return errors.WithStack(err)
		}
	}
	sm.lock.RUnlock() // avoid defer()
	return nil
}

// GetByUID returns the `ProtobufSchema` associated with the specified `schemaUID`.
//
// This is thread-safe.
func (sm *SchemaMap) GetByUID(schemaUID string) *ProtobufSchema {
	sm.lock.RLock()
	s := sm.schemaMap[schemaUID]
	sm.lock.RUnlock() // avoid defer()
	return s
}

// GetByFQName returns the first `ProtobufSchema`s that matches the specified
// fully-qualified name (e.g. `.google.protobuf.timestamp`).
//
// If more than one version of the same schema are stored in the `SchemaMap`,
// a fully-qualified name will naturally point to several distinct schemaUIDs.
// When this happens, `GetByFQName` will always return the first one to had
// been inserted in the map.
//
// This is thread-safe.
func (sm *SchemaMap) GetByFQName(fqName string) *ProtobufSchema {
	sm.lock.RLock()
	uids := sm.revmap[fqName]
	if len(uids) <= 0 {
		return nil
	}
	s := sm.schemaMap[uids[0]]
	sm.lock.RUnlock() // avoid defer()
	return s<commit_msg>Protoscan: avoid a dead-lock when the map is empty<commit_after>// Copyright © 2016 Zenly <<EMAIL>>.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package protein

import (
	"reflect"
	"sync"

	"go.uber.org/zap"

	"github.com/gogo/protobuf/protoc-gen-gogo/descriptor"
	"github.com/pkg/errors"
	"github.com/znly/protein/failure"
	"github.com/znly/protein/protoscan"
)

// -----------------------------------------------------------------------------

// SchemaMap is a thread-safe mapping & reverse-mapping of `ProtobufSchema`s.
//
// It atomically maintains two data-structures in parallel:
// a map of schemaUIDs to `ProtobufSchema`s, and
// a map of fully-qualified schema names to schemaUIDs.
//
// The `SchemaMap` is the main data-structure behind a `Transcoder`, used to
// store and retrieve every `ProtobufSchema`s that have been cached locally.
type SchemaMap struct {
	lock *sync.RWMutex
	// map[schemaUID]ProtobufSchema
	schemaMap map[string]*ProtobufSchema
	// reverse-mapping of fully-qualified names to UIDs
	// map[FQNname][]{schemaUID}
	revmap map[string][]string
}

// NewSchemaMap returns a new SchemaMap with its maps & locks pre-allocated.
func NewSchemaMap() *SchemaMap {
	return &SchemaMap{
		lock:      &sync.RWMutex{},
		schemaMap: make(map[string]*ProtobufSchema, 0),
		revmap:    make(map[string][]string, 0),
	}
}

// Add walks over the given map of `schemas` and add them to the `SchemaMap`
// while making sure to atomically maintain both the internal map and reverse-map.
func (sm *SchemaMap) Add(schemas map[string]*ProtobufSchema) *SchemaMap {
	sm.lock.Lock()
	for _, s := range schemas {
		zap.L().Debug("schema found",
			zap.String("uid", s.SchemaUID), zap.String("fq-name", s.FQName),
		)
		sm.schemaMap[s.SchemaUID] = s
		sm.revmap[s.FQName] = append(sm.revmap[s.FQName], s.SchemaUID)
	}
	sm.lock.Unlock() // avoid defer()
	return sm
}

// Size returns the number of schemas stored in the `SchemaMap`.
func (sm *SchemaMap) Size() int {
	sm.lock.RLock()
	sz := len(sm.schemaMap)
	sm.lock.RUnlock() // avoid defer()
	return sz
}

// ForEach applies the specified function `f` to every entry in the `SchemaMap`.
//
// ForEach is guaranteed to see a consistent view of the internal mapping.
func (sm *SchemaMap) ForEach(f func(s *ProtobufSchema) error) error {
	sm.lock.RLock()
	for _, s := range sm.schemaMap {
		if err := f(s); err != nil {
			sm.lock.RUnlock() // avoid defer()
			return errors.WithStack(err)
		}
	}
	sm.lock.RUnlock() // avoid defer()
	return nil
}

// GetByUID returns the `ProtobufSchema` associated with the specified `schemaUID`.
//
// This is thread-safe.
func (sm *SchemaMap) GetByUID(schemaUID string) *ProtobufSchema {
	sm.lock.RLock()
	s := sm.schemaMap[schemaUID]
	sm.lock.RUnlock() // avoid defer()
	return s
}

// GetByFQName returns the first `ProtobufSchema`s that matches the specified
// fully-qualified name (e.g. `.google.protobuf.timestamp`).
//
// If more than one version of the same schema are stored in the `SchemaMap`,
// a fully-qualified name will naturally point to several distinct schemaUIDs.
// When this happens, `GetByFQName` will always return the first one to had
// been inserted in the map.
//
// This is thread-safe.
func (sm *SchemaMap) GetByFQName(fqName string) *ProtobufSchema {
	sm.lock.RLock()
	uids := sm.revmap[fqName]
	if len(uids) <= 0 {
		sm.lock.RUnlock() // avoid defer()
		return nil
	}
	s := sm.schemaMap[uids[0]]
	sm.lock.RUnlock() // avoid defer()
	return s
{'fraction_non_alphanumeric': 0.14448958862366684, 'fraction_numerical': 0.0030472320975114273, 'mean_word_length': 3.405480984340045, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '858603', 'n_tokens_mistral': 2585, 'n_tokens_neox': 2508, 'n_words': 844}