func()

in pkg/storage/objects/host_infos.go [312:400]


func (d *hostInfoOps) CompareAndSet(
	ctx context.Context,
	hostname string,
	hostInfoDiff common.HostInfoDiff,
	compareFields common.HostInfoDiff,
) error {
	d.lock.Lock()
	defer d.lock.Unlock()

	hostInfoObject := &HostInfoObject{
		Hostname: base.NewOptionalString(hostname),
	}

	row, err := d.store.oClient.Get(
		ctx,
		hostInfoObject,
	)
	if err != nil {
		d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
		return err
	}

	if len(row) == 0 {
		return yarpcerrors.NotFoundErrorf("host info not found")
	}

	hostInfoObject.transform(row)
	info, err := newHostInfoFromHostInfoObject(hostInfoObject)
	if err != nil {
		d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
		return err
	}

	for key, value := range compareFields {
		switch key {
		case common.StateField:
			if info.GetState() != value.(hostpb.HostState) {
				d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
				return yarpcerrors.AbortedErrorf("State field does not match")
			}
		case common.GoalStateField:
			if info.GetGoalState() != value.(hostpb.HostState) {
				d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
				return yarpcerrors.AbortedErrorf("GoalState field does not match")
			}
		case common.CurrentPoolField:
			if info.GetCurrentPool() != value.(string) {
				d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
				return yarpcerrors.AbortedErrorf("CurrentPool field does not match")
			}
		case common.DesiredPoolField:
			if info.GetDesiredPool() != value.(string) {
				d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
				return yarpcerrors.AbortedErrorf("DesiredPool field does not match")
			}
		}
	}

	hostInfoObject = &HostInfoObject{
		Hostname:   base.NewOptionalString(hostname),
		UpdateTime: time.Now(),
	}

	fieldsToUpdate := []string{"UpdateTime"}
	for field, value := range hostInfoDiff {
		switch field {
		case common.StateField:
			hostInfoObject.State = value.(hostpb.HostState).String()
		case common.GoalStateField:
			hostInfoObject.GoalState = value.(hostpb.HostState).String()
		case common.CurrentPoolField:
			hostInfoObject.CurrentPool = value.(string)
		case common.DesiredPoolField:
			hostInfoObject.DesiredPool = value.(string)
		}
		fieldsToUpdate = append(fieldsToUpdate, field)
	}

	if err := d.store.oClient.Update(
		ctx,
		hostInfoObject,
		fieldsToUpdate...); err != nil {
		d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSetFail.Inc(1)
		return err
	}

	d.store.metrics.OrmHostInfoMetrics.HostInfoCompareAndSet.Inc(1)
	return nil
}