com.staros.worker.Worker Maven / Gradle / Ivy
// Copyright 2021-present StarRocks, Inc. All rights reserved.
//
// 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
//
// https://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 com.staros.worker;
import com.staros.proto.AddShardRequest;
import com.staros.proto.RemoveShardRequest;
import com.staros.proto.WorkerInfo;
import com.staros.proto.WorkerState;
import com.staros.starlet.StarletAgent;
import com.staros.starlet.StarletAgentFactory;
import com.staros.util.Text;
import com.staros.util.Writable;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.io.DataInput;
import java.io.DataOutput;
import java.io.IOException;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
public class Worker implements Writable {
private static final Logger LOG = LogManager.getLogger(Worker.class);
private final String serviceId;
private final long groupId;
private final long workerId;
private final String ipPort;
private long startTime;
// last healthy heartbeat time, Usually it is System.currentTimeMillis() in ms
private long lastSeenTime;
private AtomicReference state;
private Map workerProperties;
// Approximate number of shards reported by Starlet
// NO Need to persistent, will be updated by worker's heartbeat
// TODO: what if allows starlet to report shards in separate batches?
private long numOfShards;
private StarletAgent starletAgent;
public Worker(String serviceId, long groupId, long workerId, String ipPort) {
this.serviceId = serviceId;
this.groupId = groupId;
this.workerId = workerId;
this.ipPort = ipPort;
this.startTime = 0;
this.state = new AtomicReference(WorkerState.DOWN);
this.workerProperties = new HashMap<>();
this.starletAgent = StarletAgentFactory.newStarletAgent();
this.numOfShards = 0;
this.starletAgent.connectWorker(this);
}
public String getServiceId() {
return serviceId;
}
public long getGroupId() {
return groupId;
}
public long getWorkerId() {
return workerId;
}
public String getIpPort() {
return ipPort;
}
public WorkerState getState() {
return state.get();
}
public Map getProperties() {
return workerProperties;
}
public long getNumOfShards() {
return numOfShards < 0 ? 0 : numOfShards;
}
public boolean setState(WorkerState state) {
if (getState() == state) {
return false;
}
this.state.set(state);
LOG.info("worker {} state set to {}.", workerId, this.state);
return true;
}
public boolean isAlive() {
return getState() == WorkerState.ON;
}
public void updateLastSeenTime(long lastSeenTime) {
if (lastSeenTime > this.lastSeenTime) {
this.lastSeenTime = lastSeenTime;
}
}
public long getLastSeenTime() {
return lastSeenTime;
}
public Pair updateInfo(long sTime, Map workerProps, long numOfShards) {
boolean needPersist = false; // set to true if any worker status changes
if (workerProps != null) {
// NOTE: this might be expensive, but worker properties are expected to be only a few, so it's fine
if (!workerProperties.equals(workerProps)) {
LOG.info("worker {} properties changed, prev:{}, now:{}.", workerId, workerProperties, workerProps);
needPersist = true;
}
workerProperties = workerProps;
}
boolean restarted = false;
if (sTime == 0) {
// parameter validation
LOG.warn("worker {} start time should not be 0 in heartbeat!", workerId);
} else {
if (startTime != sTime) {
// startTime == 0 means this is the first time to get the worker's heartbeat, don't treat it as restart.
if (startTime != 0L) {
restarted = true;
SimpleDateFormat dtFormat = new SimpleDateFormat("MM-dd HH:mm:ss");
LOG.info("detect worker {} start at {}, previous start time: {}.",
workerId, dtFormat.format(new Date(sTime)), dtFormat.format(new Date(startTime)));
}
needPersist = true;
startTime = sTime;
}
}
// update number of shards hosted on the worker
this.numOfShards = numOfShards;
return Pair.of(restarted, needPersist);
}
public WorkerInfo toProtobuf() {
WorkerInfo.Builder builder = WorkerInfo.newBuilder();
builder.setServiceId(serviceId);
builder.setGroupId(groupId);
builder.setWorkerId(workerId);
builder.setIpPort(ipPort);
builder.setWorkerState(getState());
builder.putAllWorkerProperties(workerProperties);
builder.setStartTime(startTime);
return builder.build();
}
public static Worker fromProtobuf(WorkerInfo info) {
String serviceId = info.getServiceId();
long groupId = info.getGroupId();
long workerId = info.getWorkerId();
String ipPort = info.getIpPort();
WorkerState state = info.getWorkerState();
long startTime = info.getStartTime();
Map workerProps = info.getWorkerPropertiesMap();
Worker worker = new Worker(serviceId, groupId, workerId, ipPort);
worker.setState(state);
worker.updateInfo(startTime, workerProps, 0);
// use startTime as its initial lastSeenTime
worker.updateLastSeenTime(startTime);
return worker;
}
@Override
public void write(DataOutput out) throws IOException {
byte[] bytes = toProtobuf().toByteArray();
Text.writeBytes(out, bytes);
}
public static Worker read(DataInput in) throws IOException {
byte[] bytes = Text.readBytes(in);
WorkerInfo info = WorkerInfo.parseFrom(bytes);
return Worker.fromProtobuf(info);
}
public boolean equals(Object other) {
if (this == other) {
return true;
}
if (other == null) {
return false;
}
if (this.getClass() != other.getClass()) {
return false;
}
Worker otherW = (Worker) other;
if (this.workerId == otherW.workerId && this.groupId == otherW.groupId
&& this.serviceId.equals(otherW.serviceId)) {
return true;
} else {
return false;
}
}
public boolean heartbeat() {
boolean success = starletAgent.heartbeat();
if (success) {
this.updateLastSeenTime(System.currentTimeMillis());
}
return success;
}
public boolean addShard(AddShardRequest request) {
if (!isAlive()) {
return false;
}
boolean success = starletAgent.addShard(request);
if (success) {
// may not be accurate
numOfShards += request.getShardInfoCount();
}
return success;
}
public boolean removeShard(RemoveShardRequest request) {
if (!isAlive()) {
return false;
}
boolean success = starletAgent.removeShard(request);
if (success) {
// may not be accurate
numOfShards -= request.getShardIdsCount();
}
return success;
}
// FOR TEST
public StarletAgent getStarletAgent() {
return starletAgent;
}
// FOR TEST
public void cleanStarletAgent() {
starletAgent.clean();
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy