Add a refcount to dl load() and close() calls

Signed-off-by: Kevin Klues <[email protected]>

gen/nvml/lib.go

@@ -38,9 +38,10 @@ var errLibraryAlreadyLoaded = errors.New("library already loaded")
 // This includes a reference to the underlying DynamicLibrary
 type library struct {
 	sync.Mutex
-	path  string
-	flags int
-	dl    dynamicLibrary
+	path     string
+	flags    int
+	refcount refcount
+	dl       dynamicLibrary
 }
 
 // libnvml is a global instance of the nvml library.
@@ -80,20 +81,20 @@ var newDynamicLibrary = func(path string, flags int) dynamicLibrary {
 func (l *library) load() error {
 	l.Lock()
 	defer l.Unlock()
-	if l.dl != nil {
-		return nil
-	}
 
-	dl := newDynamicLibrary(l.path, l.flags)
-	err := dl.Open()
-	if err != nil {
-		return fmt.Errorf("error opening %s: %w", l.path, err)
-	}
+	err := l.refcount.FirstIn(func() error {
+		dl := newDynamicLibrary(l.path, l.flags)
+		err := dl.Open()
+		if err != nil {
+			return fmt.Errorf("error opening %s: %w", l.path, err)
+		}
 
-	l.dl = dl
-	l.updateVersionedSymbols()
+		l.dl = dl
+		l.updateVersionedSymbols()
+		return nil
+	})
 
-	return nil
+	return err
 }
 
 // close the underlying library and ensure that the global pointer to the
@@ -103,18 +104,17 @@ func (l *library) close() error {
 	l.Lock()
 	defer l.Unlock()
 
-	if l.dl == nil {
-		return nil
-	}
-
-	err := l.dl.Close()
-	if err != nil {
-		return fmt.Errorf("error closing %s: %w", l.path, err)
-	}
+	err := l.refcount.LastOut(func() error {
+		err := l.dl.Close()
+		if err != nil {
+			return fmt.Errorf("error closing %s: %w", l.path, err)
+		}
 
-	l.dl = nil
+		l.dl = nil
+		return nil
+	})
 
-	return nil
+	return err
 }
 
 // Default all versioned APIs to v1 (to infer the types)

gen/nvml/refcount.go

@@ -0,0 +1,48 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# 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 nvml
+
+import ()
+
+type refcount int
+
+func (r *refcount) FirstIn(f func() error) error {
+	(*r)++
+	if (*r) > 1 {
+		return nil
+	}
+
+	err := f()
+	if err != nil {
+		(*r)--
+	}
+
+	return err
+}
+
+func (r *refcount) LastOut(f func() error) error {
+	if (*r) != 1 {
+		return nil
+	}
+
+	err := f()
+	if err == nil {
+		(*r)--
+	}
+
+	return err
+}

pkg/nvml/lib.go

@@ -38,9 +38,10 @@ var errLibraryAlreadyLoaded = errors.New("library already loaded")
 // This includes a reference to the underlying DynamicLibrary
 type library struct {
 	sync.Mutex
-	path  string
-	flags int
-	dl    dynamicLibrary
+	path     string
+	flags    int
+	refcount refcount
+	dl       dynamicLibrary
 }
 
 // libnvml is a global instance of the nvml library.
@@ -80,20 +81,20 @@ var newDynamicLibrary = func(path string, flags int) dynamicLibrary {
 func (l *library) load() error {
 	l.Lock()
 	defer l.Unlock()
-	if l.dl != nil {
-		return nil
-	}
 
-	dl := newDynamicLibrary(l.path, l.flags)
-	err := dl.Open()
-	if err != nil {
-		return fmt.Errorf("error opening %s: %w", l.path, err)
-	}
+	err := l.refcount.FirstIn(func() error {
+		dl := newDynamicLibrary(l.path, l.flags)
+		err := dl.Open()
+		if err != nil {
+			return fmt.Errorf("error opening %s: %w", l.path, err)
+		}
 
-	l.dl = dl
-	l.updateVersionedSymbols()
+		l.dl = dl
+		l.updateVersionedSymbols()
+		return nil
+	})
 
-	return nil
+	return err
 }
 
 // close the underlying library and ensure that the global pointer to the
@@ -103,18 +104,17 @@ func (l *library) close() error {
 	l.Lock()
 	defer l.Unlock()
 
-	if l.dl == nil {
-		return nil
-	}
-
-	err := l.dl.Close()
-	if err != nil {
-		return fmt.Errorf("error closing %s: %w", l.path, err)
-	}
+	err := l.refcount.LastOut(func() error {
+		err := l.dl.Close()
+		if err != nil {
+			return fmt.Errorf("error closing %s: %w", l.path, err)
+		}
 
-	l.dl = nil
+		l.dl = nil
+		return nil
+	})
 
-	return nil
+	return err
 }
 
 // Default all versioned APIs to v1 (to infer the types)

pkg/nvml/refcount.go

@@ -0,0 +1,48 @@
+/**
+# Copyright 2024 NVIDIA CORPORATION
+#
+# 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 nvml
+
+import ()
+
+type refcount int
+
+func (r *refcount) FirstIn(f func() error) error {
+	(*r)++
+	if (*r) > 1 {
+		return nil
+	}
+
+	err := f()
+	if err != nil {
+		(*r)--
+	}
+
+	return err
+}
+
+func (r *refcount) LastOut(f func() error) error {
+	if (*r) != 1 {
+		return nil
+	}
+
+	err := f()
+	if err == nil {
+		(*r)--
+	}
+
+	return err
+}