Make TaskList generic

Signed-off-by: Maksym Pavlenko <pavlenko.maksym@gmail.com>

runtime/nsmap.go

@@ -0,0 +1,142 @@
+/*
+   Copyright The containerd Authors.
+
+   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 runtime
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/containerd/containerd/errdefs"
+	"github.com/containerd/containerd/namespaces"
+)
+
+type object interface {
+	ID() string
+}
+
+// NSMap extends Map type with a notion of namespaces passed via Context.
+type NSMap[T object] struct {
+	mu      sync.Mutex
+	objects map[string]map[string]T
+}
+
+// NewNSMap returns a new NSMap
+func NewNSMap[T object]() *NSMap[T] {
+	return &NSMap[T]{
+		objects: make(map[string]map[string]T),
+	}
+}
+
+// Get a task
+func (m *NSMap[T]) Get(ctx context.Context, id string) (T, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	namespace, err := namespaces.NamespaceRequired(ctx)
+	var t T
+	if err != nil {
+		return t, err
+	}
+	tasks, ok := m.objects[namespace]
+	if !ok {
+		return t, errdefs.ErrNotFound
+	}
+	t, ok = tasks[id]
+	if !ok {
+		return t, errdefs.ErrNotFound
+	}
+	return t, nil
+}
+
+// GetAll objects under a namespace
+func (m *NSMap[T]) GetAll(ctx context.Context, noNS bool) ([]T, error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	var o []T
+	if noNS {
+		for ns := range m.objects {
+			for _, t := range m.objects[ns] {
+				o = append(o, t)
+			}
+		}
+		return o, nil
+	}
+	namespace, err := namespaces.NamespaceRequired(ctx)
+	if err != nil {
+		return nil, err
+	}
+	tasks, ok := m.objects[namespace]
+	if !ok {
+		return o, nil
+	}
+	for _, t := range tasks {
+		o = append(o, t)
+	}
+	return o, nil
+}
+
+// Add a task
+func (m *NSMap[T]) Add(ctx context.Context, t T) error {
+	namespace, err := namespaces.NamespaceRequired(ctx)
+	if err != nil {
+		return err
+	}
+	return m.AddWithNamespace(namespace, t)
+}
+
+// AddWithNamespace adds a task with the provided namespace
+func (m *NSMap[T]) AddWithNamespace(namespace string, t T) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	id := t.ID()
+	if _, ok := m.objects[namespace]; !ok {
+		m.objects[namespace] = make(map[string]T)
+	}
+	if _, ok := m.objects[namespace][id]; ok {
+		return fmt.Errorf("%s: %w", id, errdefs.ErrAlreadyExists)
+	}
+	m.objects[namespace][id] = t
+	return nil
+}
+
+// Delete a task
+func (m *NSMap[T]) Delete(ctx context.Context, id string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	namespace, err := namespaces.NamespaceRequired(ctx)
+	if err != nil {
+		return
+	}
+	tasks, ok := m.objects[namespace]
+	if ok {
+		delete(tasks, id)
+	}
+}
+
+func (m *NSMap[T]) IsEmpty() bool {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	for ns := range m.objects {
+		if len(m.objects[ns]) > 0 {
+			return false
+		}
+	}
+
+	return true
+}