Examples with UcxCallback org.openucx.jucx.UcxCallback used on opensource projects

Example 1 with UcxCallback

use of org.openucx.jucx.UcxCallback in project twister2 by DSC-SPIDAL.

the class TWSUCXChannel method sendMessage.

@Override
public boolean sendMessage(int id, ChannelMessage message, ChannelListener callback) {
    AtomicInteger buffersLeft = new AtomicInteger(message.getBuffers().size());
    for (DataBuffer buffer : message.getBuffers()) {
        buffer.getByteBuffer().limit(buffer.getSize());
        buffer.getByteBuffer().position(0);
        int tag = this.workerId * tagWIdOffset + message.getHeader().getEdge();
        LOG.log(Level.FINE, () -> String.format("SENDING to %d[%d] : %s, TAG[%d]", id, message.getHeader().getEdge(), buffer.getByteBuffer(), tag));
        this.endpoints.get(id).sendTaggedNonBlocking(buffer.getByteBuffer(), tag, new UcxCallback() {

            @Override
            public void onSuccess(UcpRequest request) {
                pendingSendRequests.decrementAndGet();
                if (buffersLeft.decrementAndGet() == 0) {
                    callback.onSendComplete(id, message.getHeader().getEdge(), message);
                }
            }

            @Override
            public void onError(int ucsStatus, String errorMsg) {
                // This is a catastrophic failure
                LOG.severe("UCX send request failed to worker " + id + " with status " + ucsStatus + ". Error : " + errorMsg);
                throw new Twister2RuntimeException("Send request to worker : " + id + " failed. " + errorMsg);
            }
        });
        this.pendingSendRequests.incrementAndGet();
    }
    return true;
}