Add RegisterEventForDevice

src/nvml/bindings.go

@@ -17,9 +17,17 @@ const (
 	szUUID     = C.NVML_DEVICE_UUID_BUFFER_SIZE
 	szProcs    = 32
 	szProcName = 64
+
+	XidCriticalError = C.nvmlEventTypeXidCriticalError
 )
 
 type handle struct{ dev C.nvmlDevice_t }
+type EventSet struct { set C.nvmlEventSet_t }
+type Event struct {
+	UUID  *string
+	Etype uint64
+	Edata uint64
+}
 
 func uintPtr(c C.uint) *uint {
 	i := uint(c)
@@ -52,6 +60,86 @@ func init_() error {
 	return errorString(r)
 }
 
+func NewEventSet() EventSet {
+	var set C.nvmlEventSet_t
+	C.nvmlEventSetCreate(&set)
+
+	return EventSet{set}
+}
+
+func RegisterEvent(es EventSet, event int) error {
+	n, err := deviceGetCount()
+	if err != nil {
+		return err
+	}
+
+	var i uint
+	for i = 0; i < n; i++ {
+		h, err := deviceGetHandleByIndex(i)
+		if err != nil {
+			return err
+		}
+
+		r := C.nvmlDeviceRegisterEvents(h.dev, C.ulonglong(event), es.set)
+		if r != C.NVML_SUCCESS {
+			return errorString(r)
+		}
+	}
+
+	return nil
+}
+
+func RegisterEventForDevice(es EventSet, event int, uuid string) error {
+	n, err := deviceGetCount()
+	if err != nil {
+		return err
+	}
+
+	var i uint
+	for i = 0; i < n; i++ {
+		h, err := deviceGetHandleByIndex(i)
+		if err != nil {
+			return err
+		}
+
+		duuid, err := h.deviceGetUUID()
+		if err != nil {
+			return err
+		}
+
+		if *duuid != uuid {
+			continue
+		}
+
+		r := C.nvmlDeviceRegisterEvents(h.dev, C.ulonglong(event), es.set)
+		if r != C.NVML_SUCCESS {
+			return errorString(r)
+		}
+
+		return nil
+	}
+
+	return fmt.Errorf("nvml: device not found")
+}
+
+func DeleteEventSet(es EventSet) {
+	C.nvmlEventSetFree(es.set)
+}
+
+func WaitForEvent(es EventSet, timeout uint) (Event, error) {
+	var data C.nvmlEventData_t
+
+	r := C.nvmlEventSetWait(es.set, &data, C.uint(timeout))
+	uuid, _ := handle{data.device}.deviceGetUUID()
+
+	return Event{
+			UUID:  uuid,
+			Etype: uint64(data.eventType),
+			Edata: uint64(data.eventData),
+		},
+		errorString(r)
+}
+
 func shutdown() error {
 	return errorString(C.nvmlShutdown_dl())
 }