<commit_before>
#include <algorithm>
#include <iostream>
#include <cassert>
#include <sstream>
#include <boost/format.hpp>
#include <boost/algorithm/string.hpp>
#include <glog/logging.h>

#include "../include/tools.hpp"
#include "../include/aleAgent.hpp"

namespace deepRL {

void DeepQLearner::LoadPretrainedModel(const std::string& model_bin) {
  net_->CopyTrainedLayersFrom(model_bin);
}

void DeepQLearner::Initialize() {
  // Initialize net and solver

  caffe::SolverParameter solver_param;
  caffe::ReadProtoFromTextFileOrDie(solver_param_, &solver_param);
  solver_.reset(caffe::SolverRegistry<float>::CreateSolver(solver_param));

  //current q-net
  net_ = solver_->net();

  //target q-net
  target_net_ = solver_->net();

  // Cache pointers to blobs that hold Q values
  q_values_blob_ = net_->blob_by_name("q_values");

  // Initialize dummy input data with 0
  std::fill(dummy_input_data_.begin(), dummy_input_data_.end(), 0.0);

  // Cache pointers to input layers
  frames_input_layer_ =
      boost::dynamic_pointer_cast<caffe::MemoryDataLayer<float>>(
          net_->layer_by_name("frames_input_layer"));
  assert(frames_input_layer_);
  assert(deepRL::HasBlobSize(
      *net_->blob_by_name("frames"),
      1,
      kInputFrameCount,
      kCroppedFrameSize,
      kCroppedFrameSize));

  // Cache pointers to target input layers
  target_input_layer_ =
      boost::dynamic_pointer_cast<caffe::MemoryDataLayer<float>>(
          net_->layer_by_name("target_input_layer"));
  assert(target_input_layer_);
  assert(deepRL::HasBlobSize(
      *net_->blob_by_name("target"), 1, kOutputCount, 1, 1));

  // Cache pointers to filter input layers
  filter_input_layer_ =
      boost::dynamic_pointer_cast<caffe::MemoryDataLayer<float>>(
          net_->layer_by_name("filter_input_layer"));
  assert(filter_input_layer_);
  assert(deepRL::HasBlobSize(
      *net_->blob_by_name("filter"), 1, kOutputCount, 1, 1));
}

//one inputframe mode
std::vector<float>  DeepQLearner::ForwardQvalue(const InputFrames& frames_input,const NetSp& qnet){

  std::vector<float> q_values(legal_actions_.size());
  std::array<float,kInputDataSize> frames_input_data;
  for (auto j = 0; j < kInputFrameCount; ++j) {//
    const auto& frame_data = frames_input[j];
    std::copy(
        frame_data->begin(),
        frame_data->end(),
        frames_input_data.begin() + j*kCroppedFrameDataSize);
  }
  FillData2Layers(frames_input_data, dummy_input_data_, dummy_input_data_);

  qnet->Forward();//test ,now we got the qvalue

  // Get the Q values from the net
  const auto action_evaluator = [&](Action action) {
    const auto q = q_values_blob_->data_at(0, static_cast<int>(action), 0, 0);
    //std::cout << q << std::endl;
    assert(!std::isnan(q));
    return q;
  };

  std::transform(
      legal_actions_.begin(),
      legal_actions_.end(),
      q_values.begin(),
      action_evaluator);
  //std::cout << PrintQValues(q_values,legal_actions_) << std::endl;
  return q_values;
}

ActionPair DeepQLearner::MaxActionQvalue(std::vector<float> q_values){
    const auto max_idx =
        std::distance(
            q_values.begin(),
            std::max_element(q_values.begin(), q_values.end()));

    return std::make_pair(legal_actions_[max_idx],q_values[max_idx]);
}

Action DeepQLearner::SelectAction(const InputFrames& last_frames) {
  numSteps_++;
  assert(epsilon_ >= 0.0 && epsilon_ <= 1.0);
  Action action;
  double valid_epsilon_ = epsilon_;//default TRAIN
  if ( evaluate )  valid_epsilon_ = eval_epsilon;

  if (std::uniform_real_distribution<>(0.0, 1.0)(random_engine) < valid_epsilon_) {//random
    const auto random_idx =
        std::uniform_int_distribution<int>(0, legal_actions_.size() - 1)(random_engine);
    action = legal_actions_[random_idx];
    //std::cout << action_to_string(action) << " (random)";
  } else {//max greedy
    action =  MaxActionQvalue(ForwardQvalue(input_frames,net_)).first;// max
    //std::cout << action_to_string(action) << " (greedy)";
  }
  return action;
}

void DeepQLearner::StepUpdate(const Transition& tr){//only using single inputframe
 if (current_iter_%100 ==0)
      std::cout << "iteration: " << current_iter_ << " , epsilon: " <<  epsilon_ << std::endl;
   current_iter_++;

   //extract <s,a,r,s'>
   InputFrames s = std::get<0>(tr);
   const auto a = std::get<1>(tr);
   const auto r = std::get<2>(tr); assert(r >= -1.0 && r <= 1.0);
   InputFrames s_;
   const auto terminal = std::get<3>(tr) ? false : true;
   if(!terminal) {
    for (auto i = 0; i < kInputFrameCount - 1; ++i) {
      s_[i] = std::get<0>(tr)[i + 1];
    }
    s_[kInputFrameCount - 1] = std::get<3>(tr).get();
  }

  const auto target_q = terminal ?//r+maxQ'(s,a)
          r:
          r + gamma * MaxActionQvalue(ForwardQvalue(s_,target_net_)).second;
  assert(!std::isnan(target_q));

  //do data feeding -> train Q-network
  FramesLayerInputData frames_input;
  TargetLayerInputData target_input;
  FilterLayerInputData filter_input;
  std::fill(target_input.begin(), target_input.end(), 0.0f);
  std::fill(filter_input.begin(), filter_input.end(), 0.0f);

  // only the changed q has the loss
  target_input[static_cast<int>(a)] = target_q;
  filter_input[static_cast<int>(a)] = 1;

    //VLOG(1) << "filter:" << action_to_string(action) << " target:" << target;

    for (auto j = 0; j < kInputFrameCount; ++j) {
      const auto& frame_data = s[j];
      std::copy(
          frame_data->begin(),
          frame_data->end(),
          frames_input.begin() + j * kCroppedFrameDataSize);
    }

  FillData2Layers(frames_input, target_input, filter_input);

  solver_->Step(1);

  //update target_net_ with frequency param: update_frequency_
  if ( target_q_freq >0 and numSteps_%target_q_freq ==0){
        target_net_ = solver_->net(); // can we do like this ?
  }

  //epsilon decay
    //std::cout << " epsilon:" << epsilon_ << std::endl;
  if (epsilon_decay_ > 0.0) {
    if (epsilon_ != epsilon_min_){
      epsilon_ = epsilon_*epsilon_decay_;
      if (current_iter_ > epsilon_explore_idx_)
          epsilon_ = epsilon_ < epsilon_min_ ? epsilon_min_ : epsilon_;
    }
  }
}

void DeepQLearner::BatchUpdate() {

  // Sample transitions from replay memory
  //std::cout << "==>begin sampling experiences" << std::endl;
  std::vector<int> transitions = replay_memory_.sampleTransition();

  CHECK(transitions.size() == kMinibatchSize) << "Exeperience is not sampled enough";

  for (auto i = 0; i < kMinibatchSize; ++i) {

    const auto& transition = replay_memory_.getTransitionByIdx(transitions[i]);
    StepUpdate(transition);

  }

}

void DeepQLearner::FillData2Layers(
      const FramesLayerInputData& frames_input,
      const TargetLayerInputData& target_input,
      const FilterLayerInputData& filter_input) {

    frames_input_layer_->Reset(
      const_cast<float*>(frames_input.data()),//data
      dummy_input_data_.data(),//label
      1);//size n

    target_input_layer_->Reset(
      const_cast<float*>(target_input.data()),
      dummy_input_data_.data(),
      1);

    filter_input_layer_->Reset(
      const_cast<float*>(filter_input.data()),
      dummy_input_data_.data(),
      1);
}

}

<commit_msg>target_net_ update based on  numSteps_<commit_after>
#include <algorithm>
#include <iostream>
#include <cassert>
#include <sstream>
#include <boost/format.hpp>
#include <boost/algorithm/string.hpp>
#include <glog/logging.h>

#include "../include/tools.hpp"
#include "../include/aleAgent.hpp"

namespace deepRL {

void DeepQLearner::LoadPretrainedModel(const std::string& model_bin) {
  net_->CopyTrainedLayersFrom(model_bin);
}

void DeepQLearner::Initialize() {
  // Initialize net and solver

  caffe::SolverParameter solver_param;
  caffe::ReadProtoFromTextFileOrDie(solver_param_, &solver_param);
  solver_.reset(caffe::SolverRegistry<float>::CreateSolver(solver_param));

  //current q-net
  net_ = solver_->net();

  //target q-net
  target_net_ = solver_->net();

  // Cache pointers to blobs that hold Q values
  q_values_blob_ = net_->blob_by_name("q_values");

  // Initialize dummy input data with 0
  std::fill(dummy_input_data_.begin(), dummy_input_data_.end(), 0.0);

  // Cache pointers to input layers
  frames_input_layer_ =
      boost::dynamic_pointer_cast<caffe::MemoryDataLayer<float>>(
          net_->layer_by_name("frames_input_layer"));
  assert(frames_input_layer_);
  assert(deepRL::HasBlobSize(
      *net_->blob_by_name("frames"),
      1,
      kInputFrameCount,
      kCroppedFrameSize,
      kCroppedFrameSize));

  // Cache pointers to target input layers
  target_input_layer_ =
      boost::dynamic_pointer_cast<caffe::MemoryDataLayer<float>>(
          net_->layer_by_name("target_input_layer"));
  assert(target_input_layer_);
  assert(deepRL::HasBlobSize(
      *net_->blob_by_name("target"), 1, kOutputCount, 1, 1));

  // Cache pointers to filter input layers
  filter_input_layer_ =
      boost::dynamic_pointer_cast<caffe::MemoryDataLayer<float>>(
          net_->layer_by_name("filter_input_layer"));
  assert(filter_input_layer_);
  assert(deepRL::HasBlobSize(
      *net_->blob_by_name("filter"), 1, kOutputCount, 1, 1));
}

//one inputframe mode
std::vector<float>  DeepQLearner::ForwardQvalue(const InputFrames& frames_input,const NetSp& qnet){

  std::vector<float> q_values(legal_actions_.size());
  std::array<float,kInputDataSize> frames_input_data;
  for (auto j = 0; j < kInputFrameCount; ++j) {//
    const auto& frame_data = frames_input[j];
    std::copy(
        frame_data->begin(),
        frame_data->end(),
        frames_input_data.begin() + j*kCroppedFrameDataSize);
  }
  FillData2Layers(frames_input_data, dummy_input_data_, dummy_input_data_);

  qnet->Forward();//test ,now we got the qvalue

  // Get the Q values from the net
  const auto action_evaluator = [&](Action action) {
    const auto q = q_values_blob_->data_at(0, static_cast<int>(action), 0, 0);
    //std::cout << q << std::endl;
    assert(!std::isnan(q));
    return q;
  };

  std::transform(
      legal_actions_.begin(),
      legal_actions_.end(),
      q_values.begin(),
      action_evaluator);
  //std::cout << PrintQValues(q_values,legal_actions_) << std::endl;
  return q_values;
}

ActionPair DeepQLearner::MaxActionQvalue(std::vector<float> q_values){
    const auto max_idx =
        std::distance(
            q_values.begin(),
            std::max_element(q_values.begin(), q_values.end()));

    return std::make_pair(legal_actions_[max_idx],q_values[max_idx]);
}

Action DeepQLearner::SelectAction(const InputFrames& last_frames) {
  numSteps_++;
  //update target_net_ with frequency param: target_q_freq
  if ( target_q_freq >0 and numSteps_%target_q_freq ==0){
        target_net_ = solver_->net(); // can we do like this ?
  }
  assert(epsilon_ >= 0.0 && epsilon_ <= 1.0);
  Action action;
  double valid_epsilon_ = epsilon_;//default TRAIN
  if ( evaluate )  valid_epsilon_ = eval_epsilon;
  if (std::uniform_real_distribution<>(0.0, 1.0)(random_engine) < valid_epsilon_) {//random
    const auto random_idx =
        std::uniform_int_distribution<int>(0, legal_actions_.size() - 1)(random_engine);
    action = legal_actions_[random_idx];
    //std::cout << action_to_string(action) << " (random)";
  } else {//max greedy
    action =  MaxActionQvalue(ForwardQvalue(input_frames,net_)).first;// max
    //std::cout << action_to_string(action) << " (greedy)";
  }
  return action;
}

void DeepQLearner::StepUpdate(const Transition& tr){//only using single inputframe
 if (current_iter_%100 ==0)
      std::cout << "iteration: " << current_iter_ << " , epsilon: " <<  epsilon_ << std::endl;
   current_iter_++;

   //extract <s,a,r,s'>
   InputFrames s = std::get<0>(tr);
   const auto a = std::get<1>(tr);
   const auto r = std::get<2>(tr); assert(r >= -1.0 && r <= 1.0);
   InputFrames s_;
   const auto terminal = std::get<3>(tr) ? false : true;
   if(!terminal) {
    for (auto i = 0; i < kInputFrameCount - 1; ++i) {
      s_[i] = std::get<0>(tr)[i + 1];
    }
    s_[kInputFrameCount - 1] = std::get<3>(tr).get();
  }

  const auto target_q = terminal ?//r+maxQ'(s,a)
          r:
          r + gamma * MaxActionQvalue(ForwardQvalue(s_,target_net_)).second;
  assert(!std::isnan(target_q));

  //do data feeding -> train Q-network
  FramesLayerInputData frames_input;
  TargetLayerInputData target_input;
  FilterLayerInputData filter_input;
  std::fill(target_input.begin(), target_input.end(), 0.0f);
  std::fill(filter_input.begin(), filter_input.end(), 0.0f);

  // only the changed q has the loss
  target_input[static_cast<int>(a)] = target_q;
  filter_input[static_cast<int>(a)] = 1;

    //VLOG(1) << "filter:" << action_to_string(action) << " target:" << target;

    for (auto j = 0; j < kInputFrameCount; ++j) {
      const auto& frame_data = s[j];
      std::copy(
          frame_data->begin(),
          frame_data->end(),
          frames_input.begin() + j * kCroppedFrameDataSize);
    }

  FillData2Layers(frames_input, target_input, filter_input);

  solver_->Step(1);

  //epsilon decay
    //std::cout << " epsilon:" << epsilon_ << std::endl;
  if (epsilon_decay_ > 0.0) {
    if (epsilon_ != epsilon_min_){
      epsilon_ = epsilon_*epsilon_decay_;
      if (current_iter_ > epsilon_explore_idx_)
          epsilon_ = epsilon_ < epsilon_min_ ? epsilon_min_ : epsilon_;
    }
  }
}

void DeepQLearner::BatchUpdate() {

  // Sample transitions from replay memory
  //std::cout << "==>begin sampling experiences" << std::endl;
  std::vector<int> transitions = replay_memory_.sampleTransition();

  CHECK(transitions.size() == kMinibatchSize) << "Exeperience is not sampled enough";

  for (auto i = 0; i < kMinibatchSize; ++i) {

    const auto& transition = replay_memory_.getTransitionByIdx(transitions[i]);
    StepUpdate(transition);

  }

}

void DeepQLearner::FillData2Layers(
      const FramesLayerInputData& frames_input,
      const TargetLayerInputData& target_input,
      const FilterLayerInputData& filter_input) {

    frames_input_layer_->Reset(
      const_cast<float*>(frames_input.data()),//data
      dummy_input_data_.data(),//label
      1);//size n

    target_input_layer_->Reset(
      const_cast<float*>(target_input.data()),
      dummy_input_data_.data(),
      1);

    filter_input_layer_->Reset(
      const_cast<float*>(filter_input.data()),
      dummy_input_data_.data(),
      1);
}

}

<|endoftext|>


<commit_before>#include "main.h"

// --- Params --- //

NoiseParams noise_params;
WaterParams water_params;
LightParams light_params;
TextureParams texture_params;
WindowParams window_params{ WIDTH, HEIGHT };
GridParams grid_params;

// --- Scene Objects --- //

Grid grid;
Water water;
Box box;
Sky sky;

// --- Buffers --- //

Framebuffer fbw(WIDTH, HEIGHT);
Depthbuffer fb_water_depth(WIDTH, HEIGHT);

// --- Textures --- //

GLuint tex_height;
GLuint tex_mirror;
GLuint tex_water_depth;
GLuint tex_normal_map;

// --- Other --- //

NoiseGenerator noise_generator(&tex_height, &noise_params);
Camera camera(&window_params);

// --- TODO: refactor and move bezier related code to a better location */

NAVIGATION_MODE navmode;

Bezier bezier;

// Gets called when the windows is resized.
void resize_callback(int width, int height) {
    window_params.width = width;
	window_params.height = height;

    glViewport(0, 0, width, height);
	fbw.resize(width, height);
	fb_water_depth.resize(width, height);

	camera.compute_projection_matrix();
}

void compute_height_map() {
	noise_generator.renderFractal();
	box.set_height_texture(tex_height);
	grid.set_height_texture(tex_height);
	water.set_height_texture(tex_height);
	water.set_mirror_texture(tex_mirror);
	water.set_normal_map(tex_normal_map);
}

void init(){

    // Sets background color.
    glClearColor(/*gray*/ .937,.937,.937, /*solid*/1.0);

    // Enable depth test.
    glEnable(GL_DEPTH_TEST);

	bezier.init();

	initParams();
	initTextures();
	initSceneObjects();

	noise_generator.init();

	navmode = FREE;

	compute_height_map();

#ifdef WITH_ANTTWEAKBAR

	initAntTwBar(&grid_params, &window_params, &noise_params, &water_params, &texture_params);
#endif
	check_error_gl();
}

// Gets called for every frame.
void display(){
	opengp::update_title_fps("Terrain");

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	if (navmode == FREE) {
		// Render the water reflect
		fbw.bind();
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		grid.draw(camera.get_view_matrix_mirrored(), camera.get_projection_matrix(), true);
		fbw.unbind();
		glViewport(0, 0, window_params.width, window_params.height);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		fb_water_depth.bind();
		glClear(GL_DEPTH_BUFFER_BIT | GL_COLOR_BUFFER_BIT);
		grid.draw(camera.get_view_matrix(), camera.get_projection_matrix(), false);
		fb_water_depth.unbind();

		grid.draw(camera.get_view_matrix(), camera.get_projection_matrix(), false);
		water.draw(camera.get_view_matrix(), camera.get_projection_matrix());
		box.draw(camera.get_view_matrix(), camera.get_projection_matrix());
		sky.draw(camera.get_view_matrix(), camera.get_projection_matrix());

		camera.move();
	}

	if (navmode == BEZIER) {
		// TODO
		vec3 cam_pos(2.0f, 2.0f, 2.0f);
		vec3 cam_look(0.0f, 0.0f, 0.0f);
		vec3 cam_up(0.0f, 0.0f, 1.0f);

		float t = (glfwGetTime() - bezier.get_start_time()) / bezier.get_travel_time();

		if (t >= 1) {
			bezier.set_start_time(glfwGetTime());
			t = t - 1;
		}

		bezier.pos_curve_sample_point(t, cam_pos);
		bezier.cam_look_sample_point(t, cam_look);

		mat4 view_bezier = Eigen::lookAt(cam_pos, cam_look, cam_up);

		// Render the water reflect
		fbw.bind();
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		grid.draw(camera.get_view_matrix_mirrored(), camera.get_projection_matrix(), true);
		fbw.unbind();
		glViewport(0, 0, window_params.width, window_params.height);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		fb_water_depth.bind();
		glClear(GL_DEPTH_BUFFER_BIT | GL_COLOR_BUFFER_BIT);
		grid.draw(view_bezier, camera.get_projection_matrix(), false);
		fb_water_depth.unbind();

		grid.draw(view_bezier, camera.get_projection_matrix(), false);
		water.draw(view_bezier, camera.get_projection_matrix());
		box.draw(view_bezier, camera.get_projection_matrix());
		sky.draw(view_bezier, camera.get_projection_matrix());
	}

#ifdef WITH_ANTTWEAKBAR
	TwDraw();
#endif
}

void cleanup(){
#ifdef WITH_ANTTWEAKBAR
    TwTerminate();
#endif
	bezier.cleanup();
}

int main(int, char**){
    glfwInitWindowSize(window_params.width, window_params.height);
    glfwCreateWindow("Terrain");
    glfwDisplayFunc(display);

	glfwSetWindowSizeCallback(OnWindowSize);
	glfwSetMouseButtonCallback(OnMouseButton);
	glfwSetMousePosCallback(OnMousePos);
	glfwSetMouseWheelCallback(OnMouseWheel);
	glfwSetKeyCallback(OnKey);
	glfwSetCharCallback(OnChar);

    init();
    glfwMainLoop();
    cleanup();
    return EXIT_SUCCESS;
}

// --- Helpers --- //

void initParams() {

	// --- Window ---
	window_params.width = WIDTH;
	window_params.height = HEIGHT;

	// --- Grid ---
	grid_params.resolution = 100;
	grid_params.length = 2;
	grid_params.width = 2;

	// --- Noise ---
	noise_params.noise_type = PERLIN_NOISE;
	noise_params.fractal_type = FBM;
	noise_params.noise_effect = NO_EFFECT;
	noise_params.fractal_effect = NO_EFFECT;
	noise_params.height = 1;
	noise_params.width = 1;
	noise_params.offset = 0.0f;
	noise_params.amplitude = 1.0f;
	noise_params.H = 1.2f;
	noise_params.lacunarity = 2;
	noise_params.octaves = 12;
	noise_params.seed = glfwGetTime();
	noise_params.seed -= floor(noise_params.seed);

	// --- Water ---
	water_params.height = 0;
	water_params.color = vec3(0.4f, 0.55f, 0.6f);
	water_params.depth_alpha_factor = 1.4f;
	water_params.depth_color_factor = 0.05f;
	water_params.transparency = 0.4f;

	water_params.reflection_factor = 0.25f;
	/*water_params.waves_speed = 0.04f;
	water_params.waves_tile_factor = 2.0;
	water_params.waves_amplitude = 1.0f;*/

	// --- Light ---
	light_params.Ia = vec3(0.7f, 0.7f, 0.7f);
	light_params.Id = vec3(0.3f, 0.3f, 0.3f);
	light_params.position = vec3(2.0f, 2.0f, 2.0f);

	// --- Texture ---
	texture_params.sand_min_height = -1;
	texture_params.sand_max_height = 0.15f;
	texture_params.grass_max_height = 1.0f;
	texture_params.sand_max_slope = 0.92f;
	texture_params.grass_max_slope = 0.9f;

	texture_params.grass_h_transition = 3;
	texture_params.grass_s_transition = 50;
	texture_params.sand_h_transition = 5;
	texture_params.sand_s_transition = 50;
}

void initSceneObjects() {
	box.init(&grid_params, &water_params, &noise_params);
	grid.init(&grid_params, &light_params, &texture_params);
	tex_mirror = fbw.init_texture();
	fbw.init(GL_RGB8, GL_RGB, GL_UNSIGNED_BYTE);
	tex_water_depth = fb_water_depth.init_texture();
	fb_water_depth.init();

	water.init(&grid_params, &water_params, &light_params);
	water.set_depth_texture(tex_water_depth);
	sky.init();

	check_error_gl();
}

void initTextures() {
	glGenTextures(1, &tex_normal_map);
	glBindTexture(GL_TEXTURE_2D, tex_normal_map);
	glfwLoadTexture2D("textures/water/tex_normal_map.tga", 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
}

// --- Callbacks --- //

// Callback function called by GLFW when a mouse button is clicked
void GLFWCALL OnMouseButton(int glfwButton, int glfwAction)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventMouseButtonGLFW(glfwButton, glfwAction))   // Send event to AntTweakBar
	{
#endif
		if (navmode == BEZIER) {
			std::cout << "[Warning] No mouse interaction in Bezier mode.\n" << std::flush;
			return;
		}

		// Event not handled by AntTweakBar, we handle it ourselves
		camera.mouse_button(glfwButton, glfwAction);
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW when mouse has moved
void GLFWCALL OnMousePos(int mouseX, int mouseY)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventMousePosGLFW(mouseX, mouseY))  // Send event to AntTweakBar
	{
#endif
		if (navmode == BEZIER) {
			return;
		}

		camera.mouse_pos(mouseX, mouseY);
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW on mouse wheel event
void GLFWCALL OnMouseWheel(int pos)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventMouseWheelGLFW(pos))   // Send event to AntTweakBar
	{
#endif
		// Nothing for the moment
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW on key event
void GLFWCALL OnKey(int glfwKey, int glfwAction)
{

#ifdef WITH_ANTTWEAKBAR
	if (!TwEventKeyGLFW(glfwKey, glfwAction))  // Send event to AntTweakBar
	{
#endif

		if (glfwKey == GLFW_KEY_ESC && glfwAction == GLFW_PRESS) // Want to quit?
			glfwCloseWindow();
		else
		{
			camera.keyboard(glfwKey, glfwAction);

			if (glfwAction != GLFW_RELEASE) return; ///< only act on release

			switch (glfwKey){
			case '1':
				navmode = FREE;
				std::cout << "[Info] Running in free mode\n" << std::flush;
				isVisible(true);
				break;
			case '2':
				navmode = BEZIER;
				std::cout << "[Info] Running in bezier mode\n" << std::flush;
				bezier.set_start_time(glfwGetTime());
				isVisible(false);
				break;
			default:
				std::cout << "[Warning] No keyboard interaction in bezier mode.\n" << std::flush;
				break;
			}
		}
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW on char event
void GLFWCALL OnChar(int glfwChar, int glfwAction)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventCharGLFW(glfwChar, glfwAction))    // Send event to AntTweakBar
	{
#endif
		// Nothing for the moment
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW when window size changes
void GLFWCALL OnWindowSize(int width, int height)
{
#ifdef WITH_ANTTWEAKBAR
	// Send the new window size to AntTweakBar
	TwWindowSize(width, height);
#endif

	resize_callback(width, height);
}<commit_msg>Cleaner display()<commit_after>#include "main.h"

// --- Params --- //

NoiseParams noise_params;
WaterParams water_params;
LightParams light_params;
TextureParams texture_params;
WindowParams window_params{ WIDTH, HEIGHT };
GridParams grid_params;

// --- Scene Objects --- //

Grid grid;
Water water;
Box box;
Sky sky;

// --- Buffers --- //

Framebuffer fbw(WIDTH, HEIGHT);
Depthbuffer fb_water_depth(WIDTH, HEIGHT);

// --- Textures --- //

GLuint tex_height;
GLuint tex_mirror;
GLuint tex_water_depth;
GLuint tex_normal_map;

// --- Other --- //

NoiseGenerator noise_generator(&tex_height, &noise_params);
Camera camera(&window_params);

// --- TODO: refactor and move bezier related code to a better location */

NAVIGATION_MODE navmode;

Bezier bezier;

// Gets called when the windows is resized.
void resize_callback(int width, int height) {
    window_params.width = width;
	window_params.height = height;

    glViewport(0, 0, width, height);
	fbw.resize(width, height);
	fb_water_depth.resize(width, height);

	camera.compute_projection_matrix();
}

void compute_height_map() {
	noise_generator.renderFractal();
	box.set_height_texture(tex_height);
	grid.set_height_texture(tex_height);
	water.set_height_texture(tex_height);
	water.set_mirror_texture(tex_mirror);
	water.set_normal_map(tex_normal_map);
}

void init(){

    // Sets background color.
    glClearColor(/*gray*/ .937,.937,.937, /*solid*/1.0);

    // Enable depth test.
    glEnable(GL_DEPTH_TEST);

	bezier.init();

	initParams();
	initTextures();
	initSceneObjects();

	noise_generator.init();

	navmode = FREE;

	compute_height_map();

#ifdef WITH_ANTTWEAKBAR

	initAntTwBar(&grid_params, &window_params, &noise_params, &water_params, &texture_params);
#endif
	check_error_gl();
}

// Gets called for every frame.
void display(){
	opengp::update_title_fps("Terrain");

    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	mat4 view_matrix = camera.get_view_matrix();
	mat4 view_matrix_mirrored = camera.get_view_matrix_mirrored(); // for water reflections

	if (navmode == BEZIER) {
		vec3 cam_pos(2.0f, 2.0f, 2.0f);
		vec3 cam_look(0.0f, 0.0f, 0.0f);
		vec3 cam_up(0.0f, 0.0f, 1.0f);

		float t = (glfwGetTime() - bezier.get_start_time()) / bezier.get_travel_time();

		if (t >= 1) {
			bezier.set_start_time(glfwGetTime());
			t = t - 1;
		}

		bezier.pos_curve_sample_point(t, cam_pos);
		bezier.cam_look_sample_point(t, cam_look);

		view_matrix = Eigen::lookAt(cam_pos, cam_look, cam_up);
	} else {
		// FREE/default mode
		camera.move();
	}

	// Render the water reflect
	fbw.bind();
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	grid.draw(view_matrix_mirrored, camera.get_projection_matrix(), true);
	fbw.unbind();
	glViewport(0, 0, window_params.width, window_params.height);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	fb_water_depth.bind();
	glClear(GL_DEPTH_BUFFER_BIT | GL_COLOR_BUFFER_BIT);
	grid.draw(view_matrix, camera.get_projection_matrix(), false);
	fb_water_depth.unbind();

	grid.draw(view_matrix, camera.get_projection_matrix(), false);
	water.draw(view_matrix, camera.get_projection_matrix());
	box.draw(view_matrix, camera.get_projection_matrix());
	sky.draw(view_matrix, camera.get_projection_matrix());

#ifdef WITH_ANTTWEAKBAR
	TwDraw();
#endif
}

void cleanup(){
#ifdef WITH_ANTTWEAKBAR
    TwTerminate();
#endif
	bezier.cleanup();
}

int main(int, char**){
    glfwInitWindowSize(window_params.width, window_params.height);
    glfwCreateWindow("Terrain");
    glfwDisplayFunc(display);

	glfwSetWindowSizeCallback(OnWindowSize);
	glfwSetMouseButtonCallback(OnMouseButton);
	glfwSetMousePosCallback(OnMousePos);
	glfwSetMouseWheelCallback(OnMouseWheel);
	glfwSetKeyCallback(OnKey);
	glfwSetCharCallback(OnChar);

    init();
    glfwMainLoop();
    cleanup();
    return EXIT_SUCCESS;
}

// --- Helpers --- //

void initParams() {

	// --- Window ---
	window_params.width = WIDTH;
	window_params.height = HEIGHT;

	// --- Grid ---
	grid_params.resolution = 100;
	grid_params.length = 2;
	grid_params.width = 2;

	// --- Noise ---
	noise_params.noise_type = PERLIN_NOISE;
	noise_params.fractal_type = FBM;
	noise_params.noise_effect = NO_EFFECT;
	noise_params.fractal_effect = NO_EFFECT;
	noise_params.height = 1;
	noise_params.width = 1;
	noise_params.offset = 0.0f;
	noise_params.amplitude = 1.0f;
	noise_params.H = 1.2f;
	noise_params.lacunarity = 2;
	noise_params.octaves = 12;
	noise_params.seed = glfwGetTime();
	noise_params.seed -= floor(noise_params.seed);

	// --- Water ---
	water_params.height = 0;
	water_params.color = vec3(0.4f, 0.55f, 0.6f);
	water_params.depth_alpha_factor = 1.4f;
	water_params.depth_color_factor = 0.05f;
	water_params.transparency = 0.4f;

	water_params.reflection_factor = 0.25f;
	/*water_params.waves_speed = 0.04f;
	water_params.waves_tile_factor = 2.0;
	water_params.waves_amplitude = 1.0f;*/

	// --- Light ---
	light_params.Ia = vec3(0.7f, 0.7f, 0.7f);
	light_params.Id = vec3(0.3f, 0.3f, 0.3f);
	light_params.position = vec3(2.0f, 2.0f, 2.0f);

	// --- Texture ---
	texture_params.sand_min_height = -1;
	texture_params.sand_max_height = 0.15f;
	texture_params.grass_max_height = 1.0f;
	texture_params.sand_max_slope = 0.92f;
	texture_params.grass_max_slope = 0.9f;

	texture_params.grass_h_transition = 3;
	texture_params.grass_s_transition = 50;
	texture_params.sand_h_transition = 5;
	texture_params.sand_s_transition = 50;
}

void initSceneObjects() {
	box.init(&grid_params, &water_params, &noise_params);
	grid.init(&grid_params, &light_params, &texture_params);
	tex_mirror = fbw.init_texture();
	fbw.init(GL_RGB8, GL_RGB, GL_UNSIGNED_BYTE);
	tex_water_depth = fb_water_depth.init_texture();
	fb_water_depth.init();

	water.init(&grid_params, &water_params, &light_params);
	water.set_depth_texture(tex_water_depth);
	sky.init();

	check_error_gl();
}

void initTextures() {
	glGenTextures(1, &tex_normal_map);
	glBindTexture(GL_TEXTURE_2D, tex_normal_map);
	glfwLoadTexture2D("textures/water/tex_normal_map.tga", 0);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR);
	glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR);
}

// --- Callbacks --- //

// Callback function called by GLFW when a mouse button is clicked
void GLFWCALL OnMouseButton(int glfwButton, int glfwAction)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventMouseButtonGLFW(glfwButton, glfwAction))   // Send event to AntTweakBar
	{
#endif
		if (navmode == BEZIER) {
			std::cout << "[Warning] No mouse interaction in Bezier mode.\n" << std::flush;
			return;
		}

		// Event not handled by AntTweakBar, we handle it ourselves
		camera.mouse_button(glfwButton, glfwAction);
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW when mouse has moved
void GLFWCALL OnMousePos(int mouseX, int mouseY)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventMousePosGLFW(mouseX, mouseY))  // Send event to AntTweakBar
	{
#endif
		if (navmode == BEZIER) {
			return;
		}

		camera.mouse_pos(mouseX, mouseY);
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW on mouse wheel event
void GLFWCALL OnMouseWheel(int pos)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventMouseWheelGLFW(pos))   // Send event to AntTweakBar
	{
#endif
		// Nothing for the moment
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW on key event
void GLFWCALL OnKey(int glfwKey, int glfwAction)
{

#ifdef WITH_ANTTWEAKBAR
	if (!TwEventKeyGLFW(glfwKey, glfwAction))  // Send event to AntTweakBar
	{
#endif

		if (glfwKey == GLFW_KEY_ESC && glfwAction == GLFW_PRESS) // Want to quit?
			glfwCloseWindow();
		else
		{
			camera.keyboard(glfwKey, glfwAction);

			if (glfwAction != GLFW_RELEASE) return; ///< only act on release

			switch (glfwKey){
			case '1':
				navmode = FREE;
				std::cout << "[Info] Running in free mode\n" << std::flush;
				isVisible(true);
				break;
			case '2':
				navmode = BEZIER;
				std::cout << "[Info] Running in bezier mode\n" << std::flush;
				bezier.set_start_time(glfwGetTime());
				isVisible(false);
				break;
			default:
				std::cout << "[Warning] No keyboard interaction in bezier mode.\n" << std::flush;
				break;
			}
		}
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW on char event
void GLFWCALL OnChar(int glfwChar, int glfwAction)
{
#ifdef WITH_ANTTWEAKBAR
	if (!TwEventCharGLFW(glfwChar, glfwAction))    // Send event to AntTweakBar
	{
#endif
		// Nothing for the moment
#ifdef WITH_ANTTWEAKBAR
	}
#endif
}

// Callback function called by GLFW when window size changes
void GLFWCALL OnWindowSize(int width, int height)
{
#ifdef WITH_ANTTWEAKBAR
	// Send the new window size to AntTweakBar
	TwWindowSize(width, height);
#endif

	resize_callback(width, height);
}<|endoftext|>


<commit_before>/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * 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.
 */

#include "perfetto/base/paged_memory.h"

#include <stdint.h>

#include "gtest/gtest.h"
#include "perfetto/base/build_config.h"
#include "src/base/test/vm_test_utils.h"

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_MACOSX) && \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_WIN) &&    \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
#include <sys/resource.h>
#endif

namespace perfetto {
namespace base {
namespace {

TEST(PagedMemoryTest, Basic) {
  const size_t kNumPages = 10;
  const size_t kSize = 4096 * kNumPages;
  void* ptr_raw = nullptr;
  {
    PagedMemory mem = PagedMemory::Allocate(kSize);
    ASSERT_TRUE(mem.IsValid());
    ASSERT_EQ(0u, reinterpret_cast<uintptr_t>(mem.Get()) % 4096);
    ptr_raw = mem.Get();
    for (size_t i = 0; i < kSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
    ASSERT_TRUE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif

#if PERFETTO_BUILDFLAG(PERFETTO_OS_LINUX) || \
    PERFETTO_BUILDFLAG(PERFETTO_OS_ANDROID)
    ASSERT_TRUE(mem.AdviseDontNeed(ptr_raw, kSize));

    // Make sure the pages were removed from the working set.
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
  }

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
  // Freed memory is necessarily not mapped in to the process.
  ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
}

TEST(PagedMemoryTest, Uncommitted) {
  constexpr size_t kNumPages = 4096;
  constexpr size_t kSize = 4096 * kNumPages;
  char* ptr_raw = nullptr;
  {
    PagedMemory mem = PagedMemory::Allocate(kSize, PagedMemory::kDontCommit);
    ASSERT_TRUE(mem.IsValid());
    ptr_raw = reinterpret_cast<char*>(mem.Get());

#if PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
    // Windows only commits the first 1024 pages.
    constexpr size_t kMappedSize = 4096 * 1024;

    for (size_t i = 0; i < kMappedSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));

    ASSERT_TRUE(vm_test_utils::IsMapped(ptr_raw, kMappedSize));

    // Next page shouldn't be mapped.
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw + kMappedSize, 4096));
    EXPECT_DEATH({ ptr_raw[kMappedSize] = 'x'; }, ".*");

    // Commit the remaining pages.
    mem.EnsureCommitted(kSize);

    for (size_t i = kMappedSize / sizeof(uint64_t);
         i < kSize / sizeof(uint64_t); i++) {
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));
    }
#elif PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
    // Fuchsia doesn't yet support paging. So this should be a no-op.
    mem.EnsureCommitted(kSize);
    for (size_t i = 0; i < kSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));
#else
    // Linux only maps on access.
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));

    // This should not have any effect.
    mem.EnsureCommitted(kSize);
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));

    for (size_t i = 0; i < kSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));
    ASSERT_TRUE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
  }

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
  // Freed memory is necessarily not mapped in to the process.
  ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
}

#if defined(ADDRESS_SANITIZER)
TEST(PagedMemoryTest, AccessUncommittedMemoryTriggersASAN) {
  EXPECT_DEATH(
      {
        constexpr size_t kNumPages = 4096;
        constexpr size_t kSize = 4096 * kNumPages;
        PagedMemory mem =
            PagedMemory::Allocate(kSize, PagedMemory::kDontCommit);
        ASSERT_TRUE(mem.IsValid());
        char* ptr_raw = reinterpret_cast<char*>(mem.Get());
        // Only the first 1024 pages are mapped.
        constexpr size_t kMappedSize = 4096 * 1024;
        ptr_raw[kMappedSize] = 'x';
        abort();
      },
      "AddressSanitizer: container-overflow.*");
}
#endif  // ADDRESS_SANITIZER

TEST(PagedMemoryTest, GuardRegions) {
  const size_t kSize = 4096;
  PagedMemory mem = PagedMemory::Allocate(kSize);
  ASSERT_TRUE(mem.IsValid());
  volatile char* raw = reinterpret_cast<char*>(mem.Get());
  EXPECT_DEATH({ raw[-1] = 'x'; }, ".*");
  EXPECT_DEATH({ raw[kSize] = 'x'; }, ".*");
}

// Disable this on:
// MacOS: because it doesn't seem to have an equivalent rlimit to bound mmap().
// Fuchsia: doesn't support rlimit.
// Sanitizers: they seem to try to shadow mmaped memory and fail due to OOMs.
#if !PERFETTO_BUILDFLAG(PERFETTO_OS_MACOSX) &&                                 \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_WIN) &&                                    \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA) && !defined(ADDRESS_SANITIZER) && \
    !defined(LEAK_SANITIZER) && !defined(THREAD_SANITIZER) &&                  \
    !defined(MEMORY_SANITIZER)
// Glibc headers hit this on RLIMIT_ macros.
#pragma GCC diagnostic push
#if defined(__clang__)
#pragma GCC diagnostic ignored "-Wdisabled-macro-expansion"
#endif
TEST(PagedMemoryTest, Unchecked) {
  const size_t kMemLimit = 256 * 1024 * 1024l;
  struct rlimit limit {
    kMemLimit, kMemLimit
  };
  // ASSERT_EXIT here is to spawn the test in a sub-process and avoid
  // propagating the setrlimit() to other test units in case of failure.
  ASSERT_EXIT(
      {
        ASSERT_EQ(0, setrlimit(RLIMIT_AS, &limit));
        auto mem = PagedMemory::Allocate(kMemLimit * 2, PagedMemory::kMayFail);
        ASSERT_FALSE(mem.IsValid());
        exit(0);
      },
      ::testing::ExitedWithCode(0), "");
}
#pragma GCC diagnostic pop
#endif

}  // namespace
}  // namespace base
}  // namespace perfetto
<commit_msg>base: Don't run destructors on subprocess exit in PagedMemoryTest.<commit_after>/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * 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.
 */

#include "perfetto/base/paged_memory.h"

#include <stdint.h>

#include "gtest/gtest.h"
#include "perfetto/base/build_config.h"
#include "src/base/test/vm_test_utils.h"

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_MACOSX) && \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_WIN) &&    \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
#include <sys/resource.h>
#endif

namespace perfetto {
namespace base {
namespace {

TEST(PagedMemoryTest, Basic) {
  const size_t kNumPages = 10;
  const size_t kSize = 4096 * kNumPages;
  void* ptr_raw = nullptr;
  {
    PagedMemory mem = PagedMemory::Allocate(kSize);
    ASSERT_TRUE(mem.IsValid());
    ASSERT_EQ(0u, reinterpret_cast<uintptr_t>(mem.Get()) % 4096);
    ptr_raw = mem.Get();
    for (size_t i = 0; i < kSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
    ASSERT_TRUE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif

#if PERFETTO_BUILDFLAG(PERFETTO_OS_LINUX) || \
    PERFETTO_BUILDFLAG(PERFETTO_OS_ANDROID)
    ASSERT_TRUE(mem.AdviseDontNeed(ptr_raw, kSize));

    // Make sure the pages were removed from the working set.
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
  }

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
  // Freed memory is necessarily not mapped in to the process.
  ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
}

TEST(PagedMemoryTest, Uncommitted) {
  constexpr size_t kNumPages = 4096;
  constexpr size_t kSize = 4096 * kNumPages;
  char* ptr_raw = nullptr;
  {
    PagedMemory mem = PagedMemory::Allocate(kSize, PagedMemory::kDontCommit);
    ASSERT_TRUE(mem.IsValid());
    ptr_raw = reinterpret_cast<char*>(mem.Get());

#if PERFETTO_BUILDFLAG(PERFETTO_OS_WIN)
    // Windows only commits the first 1024 pages.
    constexpr size_t kMappedSize = 4096 * 1024;

    for (size_t i = 0; i < kMappedSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));

    ASSERT_TRUE(vm_test_utils::IsMapped(ptr_raw, kMappedSize));

    // Next page shouldn't be mapped.
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw + kMappedSize, 4096));
    EXPECT_DEATH({ ptr_raw[kMappedSize] = 'x'; }, ".*");

    // Commit the remaining pages.
    mem.EnsureCommitted(kSize);

    for (size_t i = kMappedSize / sizeof(uint64_t);
         i < kSize / sizeof(uint64_t); i++) {
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));
    }
#elif PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
    // Fuchsia doesn't yet support paging. So this should be a no-op.
    mem.EnsureCommitted(kSize);
    for (size_t i = 0; i < kSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));
#else
    // Linux only maps on access.
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));

    // This should not have any effect.
    mem.EnsureCommitted(kSize);
    ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));

    for (size_t i = 0; i < kSize / sizeof(uint64_t); i++)
      ASSERT_EQ(0u, *(reinterpret_cast<uint64_t*>(mem.Get()) + i));
    ASSERT_TRUE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
  }

#if !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA)
  // Freed memory is necessarily not mapped in to the process.
  ASSERT_FALSE(vm_test_utils::IsMapped(ptr_raw, kSize));
#endif
}

#if defined(ADDRESS_SANITIZER)
TEST(PagedMemoryTest, AccessUncommittedMemoryTriggersASAN) {
  EXPECT_DEATH(
      {
        constexpr size_t kNumPages = 4096;
        constexpr size_t kSize = 4096 * kNumPages;
        PagedMemory mem =
            PagedMemory::Allocate(kSize, PagedMemory::kDontCommit);
        ASSERT_TRUE(mem.IsValid());
        char* ptr_raw = reinterpret_cast<char*>(mem.Get());
        // Only the first 1024 pages are mapped.
        constexpr size_t kMappedSize = 4096 * 1024;
        ptr_raw[kMappedSize] = 'x';
        abort();
      },
      "AddressSanitizer: container-overflow.*");
}
#endif  // ADDRESS_SANITIZER

TEST(PagedMemoryTest, GuardRegions) {
  const size_t kSize = 4096;
  PagedMemory mem = PagedMemory::Allocate(kSize);
  ASSERT_TRUE(mem.IsValid());
  volatile char* raw = reinterpret_cast<char*>(mem.Get());
  EXPECT_DEATH({ raw[-1] = 'x'; }, ".*");
  EXPECT_DEATH({ raw[kSize] = 'x'; }, ".*");
}

// Disable this on:
// MacOS: because it doesn't seem to have an equivalent rlimit to bound mmap().
// Fuchsia: doesn't support rlimit.
// Sanitizers: they seem to try to shadow mmaped memory and fail due to OOMs.
#if !PERFETTO_BUILDFLAG(PERFETTO_OS_MACOSX) &&                                 \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_WIN) &&                                    \
    !PERFETTO_BUILDFLAG(PERFETTO_OS_FUCHSIA) && !defined(ADDRESS_SANITIZER) && \
    !defined(LEAK_SANITIZER) && !defined(THREAD_SANITIZER) &&                  \
    !defined(MEMORY_SANITIZER)
// Glibc headers hit this on RLIMIT_ macros.
#pragma GCC diagnostic push
#if defined(__clang__)
#pragma GCC diagnostic ignored "-Wdisabled-macro-expansion"
#endif
TEST(PagedMemoryTest, Unchecked) {
  const size_t kMemLimit = 256 * 1024 * 1024l;
  struct rlimit limit {
    kMemLimit, kMemLimit
  };
  // ASSERT_EXIT here is to spawn the test in a sub-process and avoid
  // propagating the setrlimit() to other test units in case of failure.
  ASSERT_EXIT(
      {
        ASSERT_EQ(0, setrlimit(RLIMIT_AS, &limit));
        auto mem = PagedMemory::Allocate(kMemLimit * 2, PagedMemory::kMayFail);
        ASSERT_FALSE(mem.IsValid());
        // Use _exit() instead of exit() to avoid calling destructors on child
        // process death, which may interfere with the parent process's test
        // launcher expectations.
        _exit(0);
      },
      ::testing::ExitedWithCode(0), "");
}
#pragma GCC diagnostic pop
#endif

}  // namespace
}  // namespace base
}  // namespace perfetto
<|endoftext|>


<commit_before>#pragma once

#include <cstdint>
#include <string>

#include <boost/asio.hpp>

#include "backend.hpp"
#include "blackhole/utils/unique.hpp"
#include "connect.hpp"

namespace blackhole {

namespace sink {

namespace socket {

template<>
class boost_backend_t<boost::asio::ip::udp> {
    typedef boost::asio::ip::udp Protocol;

    const std::string m_host;

    boost::asio::io_service m_io_service;
    std::unique_ptr<Protocol::socket> m_socket;

public:
    static const char* name() {
        return "udp";
    }

    boost_backend_t(const std::string& host, std::uint16_t port) :
        m_host(host),
        m_socket(initialize(m_io_service, host, port))
    {
    }

    ssize_t write(const std::string& message) {
        return m_socket->send(boost::asio::buffer(message.data(), message.size()));
    }

private:
    static inline
    std::unique_ptr<Protocol::socket>
    initialize(boost::asio::io_service& io_service, const std::string& host, std::uint16_t port) {
        auto socket = utils::make_unique<Protocol::socket>(io_service);
        connect<Protocol>(io_service, *socket.get(), host, port);
        return socket;
    }
};

} // namespace socket

} // namespace sink

} // namespace blackhole
<commit_msg>[Code Style] Reordering.<commit_after>#pragma once

#include <cstdint>
#include <string>

#include <boost/asio.hpp>

#include "backend.hpp"
#include "blackhole/utils/unique.hpp"
#include "connect.hpp"

namespace blackhole {

namespace sink {

namespace socket {

template<>
class boost_backend_t<boost::asio::ip::udp> {
    typedef boost::asio::ip::udp Protocol;

    const std::string m_host;

    boost::asio::io_service m_io_service;
    std::unique_ptr<Protocol::socket> m_socket;

public:
    boost_backend_t(const std::string& host, std::uint16_t port) :
        m_host(host),
        m_socket(initialize(m_io_service, host, port))
    {
    }

    static const char* name() {
        return "udp";
    }

    ssize_t write(const std::string& message) {
        return m_socket->send(boost::asio::buffer(message.data(), message.size()));
    }

private:
    static inline
    std::unique_ptr<Protocol::socket>
    initialize(boost::asio::io_service& io_service, const std::string& host, std::uint16_t port) {
        auto socket = utils::make_unique<Protocol::socket>(io_service);
        connect<Protocol>(io_service, *socket.get(), host, port);
        return socket;
    }
};

} // namespace socket

} // namespace sink

} // namespace blackhole
<|endoftext|>


<commit_before>/* bzflag
 * Copyright (c) 1993 - 2004 Tim Riker
 *
 * This package is free software;  you can redistribute it and/or
 * modify it under the terms of the license found in the file
 * named LICENSE that should have accompanied this file.
 *
 * THIS PACKAGE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 */

// Provide BZFS with a list server connection

/* class header */
#include "ListServerConnection.h"

/* implementation headers */
#include <string.h>
#include <math.h>
#include "bzfio.h"
#include "version.h"
#include "TextUtils.h"
#include "protocol.h"

extern Address serverAddress;
extern PingPacket getTeamCounts();

ListServerLink::ListServerLink(std::string listServerURL, std::string publicizedAddress, std::string publicizedTitle)
{
  // parse url
  std::string protocol, hostname, pathname;
  int port = 80;
  bool useDefault = false;

  // use default if it can't be parsed
  if (!BzfNetwork::parseURL(listServerURL, protocol, hostname, port, pathname))
    useDefault = true;

  // use default if wrong protocol or invalid port
  if ((protocol != "http") || (port < 1) || (port > 65535))
    useDefault = true;

  // use default if bad address
  Address address = Address::getHostAddress(hostname.c_str());
  if (address.isAny())
    useDefault = true;

  // parse default list server URL if we need to; assume default works
  if (useDefault) {
    BzfNetwork::parseURL(DefaultListServerURL, protocol, hostname, port, pathname);
    DEBUG1("Provided list server URL (%s) is invalid.  Using default of %s.", listServerURL.c_str(), DefaultListServerURL);
  }

  // add to list
  this->address	   = address;
  this->port	   = port;
  this->pathname   = pathname;
  this->hostname   = hostname;
  this->linkSocket = NotConnected;

  this->publicizeAddress     = publicizedAddress;
  this->publicizeDescription = publicizedTitle;
  this->publicizeServer	     = true;  //if this c'tor is called, it's safe to publicize

  // schedule initial ADD message
  queueMessage(ListServerLink::ADD);
}

ListServerLink::ListServerLink()
{
  // does not create a usable link, so checks should be placed
  // in  all public member functions to ensure that nothing tries
  // to happen if publicizeServer is false
  this->linkSocket = NotConnected;
  this->publicizeServer = false;
}

ListServerLink::~ListServerLink()
{
  // now tell the list server that we're going away.  this can
  // take some time but we don't want to wait too long.  we do
  // our own multiplexing loop and wait for a maximum of 3 seconds
  // total.

  // if we aren't supposed to be publicizing, skip the whole thing
  // and don't waste 3 seconds.
  if (!publicizeServer)
    return;

  queueMessage(ListServerLink::REMOVE);
  TimeKeeper start = TimeKeeper::getCurrent();
  do {
    // compute timeout
    float waitTime = 3.0f - (TimeKeeper::getCurrent() - start);
    if (waitTime <= 0.0f)
      break;
    if (!isConnected()) //queueMessage should have connected us
      break;
    // check for list server socket connection
    int fdMax = -1;
    fd_set write_set;
    fd_set read_set;
    FD_ZERO(&write_set);
    FD_ZERO(&read_set);
    if (phase == ListServerLink::CONNECTING)
      _FD_SET(linkSocket, &write_set);
    else
      _FD_SET(linkSocket, &read_set);
    fdMax = linkSocket;

    // wait for socket to connect or timeout
    struct timeval timeout;
    timeout.tv_sec = long(floorf(waitTime));
    timeout.tv_usec = long(1.0e+6f * (waitTime - floorf(waitTime)));
    int nfound = select(fdMax + 1, (fd_set*)&read_set, (fd_set*)&write_set,
			0, &timeout);
    if (nfound == 0)
      // Time has gone, close and go
      break;
    // check for connection to list server
    if (FD_ISSET(linkSocket, &write_set))
      sendQueuedMessages();
    else if (FD_ISSET(linkSocket, &read_set))
      read();
  } while (true);

  // stop list server communication
  closeLink();
}

void ListServerLink::closeLink()
{
  if (isConnected()) {
    close(linkSocket);
    DEBUG4("Closing List Server\n");
    linkSocket = NotConnected;
  }
}

void ListServerLink::read()
{
  if (isConnected()) {
    char    buf[256];
    recv(linkSocket, buf, sizeof(buf), 0);
    closeLink();
    if (nextMessageType != ListServerLink::NONE)
      // There was a pending request arrived after we write:
      // we should redo all the stuff
      openLink();
  }
}

void ListServerLink::openLink()
{
  // start opening connection if not already doing so
  if (!isConnected()) {
    linkSocket = socket(AF_INET, SOCK_STREAM, 0);
    DEBUG4("Opening List Server\n");
    if (!isConnected()) {
      return;
    }

    // set to non-blocking for connect
    if (BzfNetwork::setNonBlocking(linkSocket) < 0) {
      closeLink();
      return;
    }

    // Make our connection come from our serverAddress in case we have
    // multiple/masked IPs so the list server can verify us.
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = serverAddress;

    // assign the address to the socket
    if (bind(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
      closeLink();
      return;
    }

    // connect.  this should fail with EINPROGRESS but check for
    // success just in case.
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port);
    addr.sin_addr   = address;
    if (connect(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
#if defined(_WIN32)
#undef EINPROGRESS
#define EINPROGRESS EWOULDBLOCK
#endif
      if (getErrno() != EINPROGRESS) {
	nerror("connecting to list server");
	// TODO should try to lookup dns name again, but we don't have it anymore
	closeLink();
      } else {
	phase = CONNECTING;
      }
    } else {
      // shouldn't arrive here. Just in case, clean
      closeLink();
    }
  }
}

void ListServerLink::queueMessage(MessageType type)
{
  // ignore if the server is not public
  if (!publicizeServer) return;

  // Open network connection only if closed
  if (!isConnected()) openLink();

  // record next message to send.
  nextMessageType = type;
}

void ListServerLink::sendQueuedMessages()
{
  if (!isConnected())
    return;

  if (nextMessageType == ListServerLink::ADD) {
    DEBUG3("Queuing ADD message to list server\n");
    addMe(getTeamCounts(), publicizeAddress, url_encode(publicizeDescription));
    lastAddTime = TimeKeeper::getCurrent();
  } else if (nextMessageType == ListServerLink::REMOVE) {
    DEBUG3("Queuing REMOVE message to list server\n");
    removeMe(publicizeAddress);
  }
}

void ListServerLink::addMe(PingPacket pingInfo,
			   std::string publicizedAddress,
			   std::string publicizedTitle)
{
  std::string msg;

  // encode ping reply as ascii hex digits plus NULL
  char gameInfo[PingPacketHexPackedSize + 1];
  pingInfo.packHex(gameInfo);

  // send ADD message (must send blank line)
  msg = string_util::format("GET %s?action=ADD&nameport=%s&version=%s&gameinfo=%s&title=%s HTTP/1.1\r\n"
    "Host: %s\r\nCache-Control: no-cache\r\n\r\n",
    pathname.c_str(), publicizedAddress.c_str(),
    getServerVersion(), gameInfo,
    publicizedTitle.c_str(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::removeMe(std::string publicizedAddress)
{
  std::string msg;
  // send REMOVE (must send blank line)
  msg = string_util::format("GET %s?action=REMOVE&nameport=%s HTTP/1.1\r\n"
    "Host: %s\r\nCache-Control: no-cache\r\n\r\n",
    pathname.c_str(),
    publicizedAddress.c_str(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::sendMessage(std::string message)
{
  const int bufsize = 4096;
  char msg[bufsize];
  strncpy(msg, message.c_str(),
          ((int) message.length() > bufsize) ? bufsize : message.length());
  if (strlen(msg) > 0) {
    DEBUG3("%s\n",msg);
    if (send(linkSocket, msg, strlen(msg), 0) == -1) {
      perror("List server send failed");
      DEBUG3("Unable to send to the list server!\n");
      closeLink();
    } else {
      nextMessageType = ListServerLink::NONE;
      phase           = ListServerLink::WRITTEN;
    }
  } else {
    closeLink();
  }
}
<commit_msg>whoopsie...that's a capital 'P'!<commit_after>/* bzflag
 * Copyright (c) 1993 - 2004 Tim Riker
 *
 * This package is free software;  you can redistribute it and/or
 * modify it under the terms of the license found in the file
 * named LICENSE that should have accompanied this file.
 *
 * THIS PACKAGE IS PROVIDED ``AS IS'' AND WITHOUT ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, WITHOUT LIMITATION, THE IMPLIED
 * WARRANTIES OF MERCHANTIBILITY AND FITNESS FOR A PARTICULAR PURPOSE.
 */

// Provide BZFS with a list server connection

/* class header */
#include "ListServerConnection.h"

/* implementation headers */
#include <string.h>
#include <math.h>
#include "bzfio.h"
#include "version.h"
#include "TextUtils.h"
#include "Protocol.h"

extern Address serverAddress;
extern PingPacket getTeamCounts();

ListServerLink::ListServerLink(std::string listServerURL, std::string publicizedAddress, std::string publicizedTitle)
{
  // parse url
  std::string protocol, hostname, pathname;
  int port = 80;
  bool useDefault = false;

  // use default if it can't be parsed
  if (!BzfNetwork::parseURL(listServerURL, protocol, hostname, port, pathname))
    useDefault = true;

  // use default if wrong protocol or invalid port
  if ((protocol != "http") || (port < 1) || (port > 65535))
    useDefault = true;

  // use default if bad address
  Address address = Address::getHostAddress(hostname.c_str());
  if (address.isAny())
    useDefault = true;

  // parse default list server URL if we need to; assume default works
  if (useDefault) {
    BzfNetwork::parseURL(DefaultListServerURL, protocol, hostname, port, pathname);
    DEBUG1("Provided list server URL (%s) is invalid.  Using default of %s.", listServerURL.c_str(), DefaultListServerURL);
  }

  // add to list
  this->address	   = address;
  this->port	   = port;
  this->pathname   = pathname;
  this->hostname   = hostname;
  this->linkSocket = NotConnected;

  this->publicizeAddress     = publicizedAddress;
  this->publicizeDescription = publicizedTitle;
  this->publicizeServer	     = true;  //if this c'tor is called, it's safe to publicize

  // schedule initial ADD message
  queueMessage(ListServerLink::ADD);
}

ListServerLink::ListServerLink()
{
  // does not create a usable link, so checks should be placed
  // in  all public member functions to ensure that nothing tries
  // to happen if publicizeServer is false
  this->linkSocket = NotConnected;
  this->publicizeServer = false;
}

ListServerLink::~ListServerLink()
{
  // now tell the list server that we're going away.  this can
  // take some time but we don't want to wait too long.  we do
  // our own multiplexing loop and wait for a maximum of 3 seconds
  // total.

  // if we aren't supposed to be publicizing, skip the whole thing
  // and don't waste 3 seconds.
  if (!publicizeServer)
    return;

  queueMessage(ListServerLink::REMOVE);
  TimeKeeper start = TimeKeeper::getCurrent();
  do {
    // compute timeout
    float waitTime = 3.0f - (TimeKeeper::getCurrent() - start);
    if (waitTime <= 0.0f)
      break;
    if (!isConnected()) //queueMessage should have connected us
      break;
    // check for list server socket connection
    int fdMax = -1;
    fd_set write_set;
    fd_set read_set;
    FD_ZERO(&write_set);
    FD_ZERO(&read_set);
    if (phase == ListServerLink::CONNECTING)
      _FD_SET(linkSocket, &write_set);
    else
      _FD_SET(linkSocket, &read_set);
    fdMax = linkSocket;

    // wait for socket to connect or timeout
    struct timeval timeout;
    timeout.tv_sec = long(floorf(waitTime));
    timeout.tv_usec = long(1.0e+6f * (waitTime - floorf(waitTime)));
    int nfound = select(fdMax + 1, (fd_set*)&read_set, (fd_set*)&write_set,
			0, &timeout);
    if (nfound == 0)
      // Time has gone, close and go
      break;
    // check for connection to list server
    if (FD_ISSET(linkSocket, &write_set))
      sendQueuedMessages();
    else if (FD_ISSET(linkSocket, &read_set))
      read();
  } while (true);

  // stop list server communication
  closeLink();
}

void ListServerLink::closeLink()
{
  if (isConnected()) {
    close(linkSocket);
    DEBUG4("Closing List Server\n");
    linkSocket = NotConnected;
  }
}

void ListServerLink::read()
{
  if (isConnected()) {
    char    buf[256];
    recv(linkSocket, buf, sizeof(buf), 0);
    closeLink();
    if (nextMessageType != ListServerLink::NONE)
      // There was a pending request arrived after we write:
      // we should redo all the stuff
      openLink();
  }
}

void ListServerLink::openLink()
{
  // start opening connection if not already doing so
  if (!isConnected()) {
    linkSocket = socket(AF_INET, SOCK_STREAM, 0);
    DEBUG4("Opening List Server\n");
    if (!isConnected()) {
      return;
    }

    // set to non-blocking for connect
    if (BzfNetwork::setNonBlocking(linkSocket) < 0) {
      closeLink();
      return;
    }

    // Make our connection come from our serverAddress in case we have
    // multiple/masked IPs so the list server can verify us.
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr = serverAddress;

    // assign the address to the socket
    if (bind(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
      closeLink();
      return;
    }

    // connect.  this should fail with EINPROGRESS but check for
    // success just in case.
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_port   = htons(port);
    addr.sin_addr   = address;
    if (connect(linkSocket, (CNCTType*)&addr, sizeof(addr)) < 0) {
#if defined(_WIN32)
#undef EINPROGRESS
#define EINPROGRESS EWOULDBLOCK
#endif
      if (getErrno() != EINPROGRESS) {
	nerror("connecting to list server");
	// TODO should try to lookup dns name again, but we don't have it anymore
	closeLink();
      } else {
	phase = CONNECTING;
      }
    } else {
      // shouldn't arrive here. Just in case, clean
      closeLink();
    }
  }
}

void ListServerLink::queueMessage(MessageType type)
{
  // ignore if the server is not public
  if (!publicizeServer) return;

  // Open network connection only if closed
  if (!isConnected()) openLink();

  // record next message to send.
  nextMessageType = type;
}

void ListServerLink::sendQueuedMessages()
{
  if (!isConnected())
    return;

  if (nextMessageType == ListServerLink::ADD) {
    DEBUG3("Queuing ADD message to list server\n");
    addMe(getTeamCounts(), publicizeAddress, url_encode(publicizeDescription));
    lastAddTime = TimeKeeper::getCurrent();
  } else if (nextMessageType == ListServerLink::REMOVE) {
    DEBUG3("Queuing REMOVE message to list server\n");
    removeMe(publicizeAddress);
  }
}

void ListServerLink::addMe(PingPacket pingInfo,
			   std::string publicizedAddress,
			   std::string publicizedTitle)
{
  std::string msg;

  // encode ping reply as ascii hex digits plus NULL
  char gameInfo[PingPacketHexPackedSize + 1];
  pingInfo.packHex(gameInfo);

  // send ADD message (must send blank line)
  msg = string_util::format("GET %s?action=ADD&nameport=%s&version=%s&gameinfo=%s&title=%s HTTP/1.1\r\n"
    "Host: %s\r\nCache-Control: no-cache\r\n\r\n",
    pathname.c_str(), publicizedAddress.c_str(),
    getServerVersion(), gameInfo,
    publicizedTitle.c_str(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::removeMe(std::string publicizedAddress)
{
  std::string msg;
  // send REMOVE (must send blank line)
  msg = string_util::format("GET %s?action=REMOVE&nameport=%s HTTP/1.1\r\n"
    "Host: %s\r\nCache-Control: no-cache\r\n\r\n",
    pathname.c_str(),
    publicizedAddress.c_str(),
    hostname.c_str());
  sendMessage(msg);
}

void ListServerLink::sendMessage(std::string message)
{
  const int bufsize = 4096;
  char msg[bufsize];
  strncpy(msg, message.c_str(),
          ((int) message.length() > bufsize) ? bufsize : message.length());
  if (strlen(msg) > 0) {
    DEBUG3("%s\n",msg);
    if (send(linkSocket, msg, strlen(msg), 0) == -1) {
      perror("List server send failed");
      DEBUG3("Unable to send to the list server!\n");
      closeLink();
    } else {
      nextMessageType = ListServerLink::NONE;
      phase           = ListServerLink::WRITTEN;
    }
  } else {
    closeLink();
  }
}
<|endoftext|>


<commit_before>//---------------------------------------------------------- -*- Mode: C++ -*-
// $Id$
//
// Created 2013/8/9
// Author: Mike Ovsiannikov
//
// Copyright 2013 Quantcast Corp.
//
// This file is part of Kosmos File System (KFS).
//
// 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.
//
//
//----------------------------------------------------------------------------

#include "ClientAuthContext.h"

#include "common/Properties.h"
#include "kfsio/NetConnection.h"
#include "common/Properties.h"
#include "common/MsgLogger.h"
#include "kfsio/SslFilter.h"
#include "krb/KrbClient.h"
#include "qcdio/qcdebug.h"

#include <boost/scoped_ptr.hpp>
#include <errno.h>

namespace KFS
{

using std::string;
using boost::scoped_ptr;

class ClientAuthContext::Impl
{
public:
    Impl()
        : mSharedFlag(false),
          mEnabledFlag(false),
          mParams(),
          mKrbClientPtr(),
          mSslCtxPtr(),
          mX509SslCtxPtr()
        {}
    ~Impl()
        {}
    int SetParameters(
        const char*       inParamsPrefixPtr,
        const Properties& inParameters,
        string*           outErrMsgPtr)
    {
        Properties::String theParamName;
        if (inParamsPrefixPtr) {
            theParamName.Append(inParamsPrefixPtr);
        }
        const size_t thePrefLen = theParamName.GetSize();
        Properties theParams(mParams);
        inParameters.copyWithPrefix(
            theParamName.GetPtr(), theParamName.GetSize(), theParams);
        const char* theNullStr         = 0;
        const char* theServeiceNamePtr = theParams.getValue(
            theParamName.Truncate(thePrefLen).Append(
            "krb5.serviceName"), theNullStr);
        theParamName.Truncate(thePrefLen).Append("krb5.");
        size_t theCurLen = theParamName.GetSize();
        const bool theKrbChangedFlag =
            theParams.getValue(theParamName.Append("forceReload"), 0) != 0 ||
            ! theParams.equalsWithPrefix(
                theParamName.Truncate(theCurLen).GetPtr(), theCurLen, mParams);
        KrbClientPtr theKrbClientPtr;
        if (theKrbChangedFlag && theServeiceNamePtr && theServeiceNamePtr[0]) {
            theKrbClientPtr.reset(new KrbClient());
            const char* const theErrMsgPtr = theKrbClientPtr->Init(
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.serviceHost"), theNullStr),
                theServeiceNamePtr,
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.keytab"), theNullStr),
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.clientName"), theNullStr),
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.initClientCache"), 0) != 0
            );
            if (theErrMsgPtr) {
                if (outErrMsgPtr) {
                    *outErrMsgPtr = theErrMsgPtr;
                }
                KFS_LOG_STREAM_ERROR <<
                    theParamName.Truncate(thePrefLen) <<
                    "krb5.* configuration error: " << theErrMsgPtr <<
                KFS_LOG_EOM;
                return -EINVAL;
            }
        }
        theParamName.Truncate(thePrefLen).Append("psk.tls.");
        theCurLen = theParamName.GetSize();
        const bool thePskSslChangedFlag =
            (theKrbChangedFlag &&
                    (mKrbClientPtr.get() != 0) !=
                    (theKrbClientPtr.get() != 0)) ||
            theParams.getValue(
                theParamName.Append("forceReload"), 0) != 0 ||
            ! theParams.equalsWithPrefix(
                theParamName.Truncate(theCurLen).GetPtr(), theCurLen, mParams);
        SslCtxPtr theSslCtxPtr;
        if (thePskSslChangedFlag && theParams.getValue(
                theParamName.Truncate(thePrefLen).Append(
                    "psk.tls.disable"), 0) == 0) {
            const bool kServerFlag  = false;
            const bool kPskOnlyFlag = true;
            string     theErrMsg;
            mSslCtxPtr.Set(SslFilter::CreateCtx(
                kServerFlag,
                kPskOnlyFlag,
                theParamName.Truncate(thePrefLen).Append("psk.tls.").GetPtr(),
                theParams,
                &theErrMsg
            ));
            if (! mSslCtxPtr.Get()) {
                KFS_LOG_STREAM_ERROR <<
                    theParamName.Truncate(thePrefLen) <<
                    "psk.tls.* configuration error: " << theErrMsg <<
                KFS_LOG_EOM;
                return false;
            }
        }
        theParamName.Truncate(thePrefLen).Append("X509.");
        theCurLen = theParamName.GetSize();
        const bool theX509ChangedFlag =
            theParams.getValue(
                theParamName.Append("forceReload"), 0) != 0 ||
            ! theParams.equalsWithPrefix(
                theParamName.Truncate(theCurLen).GetPtr(), theCurLen, mParams);
        SslCtxPtr theX509SslCtxPtr;
        if (theX509ChangedFlag) {
            const bool kServerFlag  = true;
            const bool kPskOnlyFlag = false;
            string     theErrMsg;
            mSslCtxPtr.Set(SslFilter::CreateCtx(
                kServerFlag,
                kPskOnlyFlag,
                theParamName.Truncate(thePrefLen).Append("X509.").GetPtr(),
                theParams,
                &theErrMsg
            ));
            if (! mSslCtxPtr.Get()) {
                KFS_LOG_STREAM_ERROR <<
                    theParamName.Truncate(thePrefLen) <<
                    "X509.* configuration error: " << theErrMsg <<
                KFS_LOG_EOM;
                return false;
            }
        }
        mParams.swap(theParams);
        if (theKrbChangedFlag) {
            mKrbClientPtr.swap(theKrbClientPtr);
        }
        if (thePskSslChangedFlag) {
            mSslCtxPtr.Swap(theSslCtxPtr);
        }
        if (theX509ChangedFlag) {
            mX509SslCtxPtr.Swap(theX509SslCtxPtr);
        }
        return 0;
    }
    int Request(
        int          inAuthType,
        int&         outAuthType,
        const char*& outBufPtr,
        int&         outBufLen,
        string*      outErrMsgPtr)
    {
        return 0;
    }
    int StartSsl(
        NetConnection& inNetConnection,
        const char*    inKeyIdPtr,
        const char*    inKeyDataPtr,
        int            inKeyDataSize,
        string*        outErrMsgPtr)
    {
        return 0;
    }
    int Response(
        int            inAuthType,
        bool           inUseSslFlag,
        const char*    inBufPtr,
        int            inBufLen,
        NetConnection& inNetConnection,
        string*        outErrMsgPtr)
    {
        return 0;
    }
    bool IsEnabled() const
        { return mEnabledFlag; }
    bool IsShared() const
        { return mSharedFlag; }
    void SetShared(
        bool inFlag)
    {
        mSharedFlag = inFlag;
    }
    int CheckAuthType(
        int     inAuthType,
        string* outErrMsgPtr)
    {
        return 0;
    }
private:
    typedef scoped_ptr<KrbClient> KrbClientPtr;
    typedef SslFilter::CtxPtr     SslCtxPtr;

    bool         mSharedFlag;
    bool         mEnabledFlag;
    Properties   mParams;
    KrbClientPtr mKrbClientPtr;
    SslCtxPtr    mSslCtxPtr;
    SslCtxPtr    mX509SslCtxPtr;
};

ClientAuthContext::ClientAuthContext()
    : mImpl(*(new Impl()))
    {}

ClientAuthContext::~ClientAuthContext()
{
    delete &mImpl;
}

    bool
ClientAuthContext::IsShared() const
{
    return mImpl.IsShared();
}

    void
ClientAuthContext::SetShared(
    bool inFlag)
{
    return mImpl.SetShared(inFlag);
}

    int
ClientAuthContext::CheckAuthType(
    int     inAuthType,
    string* outErrMsgPtr)
{
    return mImpl.CheckAuthType(inAuthType, outErrMsgPtr);
}

    int
ClientAuthContext::SetParameters(
    const char*       inParamsPrefixPtr,
    const Properties& inParameters,
    string*           outErrMsgPtr)
{
    return mImpl.SetParameters(inParamsPrefixPtr, inParameters, outErrMsgPtr);
}

    int
ClientAuthContext::Request(
    int          inAuthType,
    int&         outAuthType,
    const char*& outBufPtr,
    int&         outBufLen,
    string*      outErrMsgPtr)
{
    return mImpl.Request(
        inAuthType, outAuthType, outBufPtr, outBufLen, outErrMsgPtr);
}

    int
ClientAuthContext::StartSsl(
    NetConnection& inNetConnection,
    const char*    inKeyIdPtr,
    const char*    inKeyDataPtr,
    int            inKeyDataSize,
    string*        outErrMsgPtr)
{
    return mImpl.StartSsl(
        inNetConnection, inKeyIdPtr, inKeyDataPtr, inKeyDataSize, outErrMsgPtr);
}

    int
ClientAuthContext::Response(
    int            inAuthType,
    bool           inUseSslFlag,
    const char*    inBufPtr,
    int            inBufLen,
    NetConnection& inNetConnection,
    string*        outErrMsgPtr)
{
    return mImpl.Response(
        inAuthType, inUseSslFlag, inBufPtr, inBufLen, inNetConnection,
        outErrMsgPtr);
}

    bool
ClientAuthContext::IsEnabled() const
{
    return mImpl.IsEnabled();
}

} // namespace KFS
<commit_msg>Kfsio lib: fix client flag with authentication context X509 ssl.<commit_after>//---------------------------------------------------------- -*- Mode: C++ -*-
// $Id$
//
// Created 2013/8/9
// Author: Mike Ovsiannikov
//
// Copyright 2013 Quantcast Corp.
//
// This file is part of Kosmos File System (KFS).
//
// 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.
//
//
//----------------------------------------------------------------------------

#include "ClientAuthContext.h"

#include "common/Properties.h"
#include "kfsio/NetConnection.h"
#include "common/Properties.h"
#include "common/MsgLogger.h"
#include "kfsio/SslFilter.h"
#include "krb/KrbClient.h"
#include "qcdio/qcdebug.h"

#include <boost/scoped_ptr.hpp>
#include <errno.h>

namespace KFS
{

using std::string;
using boost::scoped_ptr;

class ClientAuthContext::Impl
{
public:
    Impl()
        : mSharedFlag(false),
          mEnabledFlag(false),
          mParams(),
          mKrbClientPtr(),
          mSslCtxPtr(),
          mX509SslCtxPtr()
        {}
    ~Impl()
        {}
    int SetParameters(
        const char*       inParamsPrefixPtr,
        const Properties& inParameters,
        string*           outErrMsgPtr)
    {
        Properties::String theParamName;
        if (inParamsPrefixPtr) {
            theParamName.Append(inParamsPrefixPtr);
        }
        const size_t thePrefLen = theParamName.GetSize();
        Properties theParams(mParams);
        inParameters.copyWithPrefix(
            theParamName.GetPtr(), theParamName.GetSize(), theParams);
        const char* theNullStr         = 0;
        const char* theServeiceNamePtr = theParams.getValue(
            theParamName.Truncate(thePrefLen).Append(
            "krb5.serviceName"), theNullStr);
        theParamName.Truncate(thePrefLen).Append("krb5.");
        size_t theCurLen = theParamName.GetSize();
        const bool theKrbChangedFlag =
            theParams.getValue(theParamName.Append("forceReload"), 0) != 0 ||
            ! theParams.equalsWithPrefix(
                theParamName.Truncate(theCurLen).GetPtr(), theCurLen, mParams);
        KrbClientPtr theKrbClientPtr;
        if (theKrbChangedFlag && theServeiceNamePtr && theServeiceNamePtr[0]) {
            theKrbClientPtr.reset(new KrbClient());
            const char* const theErrMsgPtr = theKrbClientPtr->Init(
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.serviceHost"), theNullStr),
                theServeiceNamePtr,
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.keytab"), theNullStr),
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.clientName"), theNullStr),
                theParams.getValue(
                    theParamName.Truncate(thePrefLen).Append(
                    "krb5.initClientCache"), 0) != 0
            );
            if (theErrMsgPtr) {
                if (outErrMsgPtr) {
                    *outErrMsgPtr = theErrMsgPtr;
                }
                KFS_LOG_STREAM_ERROR <<
                    theParamName.Truncate(thePrefLen) <<
                    "krb5.* configuration error: " << theErrMsgPtr <<
                KFS_LOG_EOM;
                return -EINVAL;
            }
        }
        theParamName.Truncate(thePrefLen).Append("psk.tls.");
        theCurLen = theParamName.GetSize();
        const bool thePskSslChangedFlag =
            (theKrbChangedFlag &&
                    (mKrbClientPtr.get() != 0) !=
                    (theKrbClientPtr.get() != 0)) ||
            theParams.getValue(
                theParamName.Append("forceReload"), 0) != 0 ||
            ! theParams.equalsWithPrefix(
                theParamName.Truncate(theCurLen).GetPtr(), theCurLen, mParams);
        SslCtxPtr theSslCtxPtr;
        if (thePskSslChangedFlag && theParams.getValue(
                theParamName.Truncate(thePrefLen).Append(
                    "psk.tls.disable"), 0) == 0) {
            const bool kServerFlag  = false;
            const bool kPskOnlyFlag = true;
            string     theErrMsg;
            mSslCtxPtr.Set(SslFilter::CreateCtx(
                kServerFlag,
                kPskOnlyFlag,
                theParamName.Truncate(thePrefLen).Append("psk.tls.").GetPtr(),
                theParams,
                &theErrMsg
            ));
            if (! mSslCtxPtr.Get()) {
                KFS_LOG_STREAM_ERROR <<
                    theParamName.Truncate(thePrefLen) <<
                    "psk.tls.* configuration error: " << theErrMsg <<
                KFS_LOG_EOM;
                return false;
            }
        }
        theParamName.Truncate(thePrefLen).Append("X509.");
        theCurLen = theParamName.GetSize();
        const bool theX509ChangedFlag =
            theParams.getValue(
                theParamName.Append("forceReload"), 0) != 0 ||
            ! theParams.equalsWithPrefix(
                theParamName.Truncate(theCurLen).GetPtr(), theCurLen, mParams);
        SslCtxPtr theX509SslCtxPtr;
        if (theX509ChangedFlag) {
            const bool kServerFlag  = false;
            const bool kPskOnlyFlag = false;
            string     theErrMsg;
            mSslCtxPtr.Set(SslFilter::CreateCtx(
                kServerFlag,
                kPskOnlyFlag,
                theParamName.Truncate(thePrefLen).Append("X509.").GetPtr(),
                theParams,
                &theErrMsg
            ));
            if (! mSslCtxPtr.Get()) {
                KFS_LOG_STREAM_ERROR <<
                    theParamName.Truncate(thePrefLen) <<
                    "X509.* configuration error: " << theErrMsg <<
                KFS_LOG_EOM;
                return false;
            }
        }
        mParams.swap(theParams);
        if (theKrbChangedFlag) {
            mKrbClientPtr.swap(theKrbClientPtr);
        }
        if (thePskSslChangedFlag) {
            mSslCtxPtr.Swap(theSslCtxPtr);
        }
        if (theX509ChangedFlag) {
            mX509SslCtxPtr.Swap(theX509SslCtxPtr);
        }
        return 0;
    }
    int Request(
        int          inAuthType,
        int&         outAuthType,
        const char*& outBufPtr,
        int&         outBufLen,
        string*      outErrMsgPtr)
    {
        return 0;
    }
    int StartSsl(
        NetConnection& inNetConnection,
        const char*    inKeyIdPtr,
        const char*    inKeyDataPtr,
        int            inKeyDataSize,
        string*        outErrMsgPtr)
    {
        return 0;
    }
    int Response(
        int            inAuthType,
        bool           inUseSslFlag,
        const char*    inBufPtr,
        int            inBufLen,
        NetConnection& inNetConnection,
        string*        outErrMsgPtr)
    {
        return 0;
    }
    bool IsEnabled() const
        { return mEnabledFlag; }
    bool IsShared() const
        { return mSharedFlag; }
    void SetShared(
        bool inFlag)
    {
        mSharedFlag = inFlag;
    }
    int CheckAuthType(
        int     inAuthType,
        string* outErrMsgPtr)
    {
        return 0;
    }
private:
    typedef scoped_ptr<KrbClient> KrbClientPtr;
    typedef SslFilter::CtxPtr     SslCtxPtr;

    bool         mSharedFlag;
    bool         mEnabledFlag;
    Properties   mParams;
    KrbClientPtr mKrbClientPtr;
    SslCtxPtr    mSslCtxPtr;
    SslCtxPtr    mX509SslCtxPtr;
};

ClientAuthContext::ClientAuthContext()
    : mImpl(*(new Impl()))
    {}

ClientAuthContext::~ClientAuthContext()
{
    delete &mImpl;
}

    bool
ClientAuthContext::IsShared() const
{
    return mImpl.IsShared();
}

    void
ClientAuthContext::SetShared(
    bool inFlag)
{
    return mImpl.SetShared(inFlag);
}

    int
ClientAuthContext::CheckAuthType(
    int     inAuthType,
    string* outErrMsgPtr)
{
    return mImpl.CheckAuthType(inAuthType, outErrMsgPtr);
}

    int
ClientAuthContext::SetParameters(
    const char*       inParamsPrefixPtr,
    const Properties& inParameters,
    string*           outErrMsgPtr)
{
    return mImpl.SetParameters(inParamsPrefixPtr, inParameters, outErrMsgPtr);
}

    int
ClientAuthContext::Request(
    int          inAuthType,
    int&         outAuthType,
    const char*& outBufPtr,
    int&         outBufLen,
    string*      outErrMsgPtr)
{
    return mImpl.Request(
        inAuthType, outAuthType, outBufPtr, outBufLen, outErrMsgPtr);
}

    int
ClientAuthContext::StartSsl(
    NetConnection& inNetConnection,
    const char*    inKeyIdPtr,
    const char*    inKeyDataPtr,
    int            inKeyDataSize,
    string*        outErrMsgPtr)
{
    return mImpl.StartSsl(
        inNetConnection, inKeyIdPtr, inKeyDataPtr, inKeyDataSize, outErrMsgPtr);
}

    int
ClientAuthContext::Response(
    int            inAuthType,
    bool           inUseSslFlag,
    const char*    inBufPtr,
    int            inBufLen,
    NetConnection& inNetConnection,
    string*        outErrMsgPtr)
{
    return mImpl.Response(
        inAuthType, inUseSslFlag, inBufPtr, inBufLen, inNetConnection,
        outErrMsgPtr);
}

    bool
ClientAuthContext::IsEnabled() const
{
    return mImpl.IsEnabled();
}

} // namespace KFS
<|endoftext|>


<commit_before>#include "resources_tracker.h"

#include <cassert>
#include <stdexcept>

namespace Meteor {

ResourcesTracker::ResourcesTracker(size_t numResources)
    : NumResources(numResources)
{
}

/// @throw runtime_error
void ResourcesTracker::ValidateResource(const Resource& resource) {
    if (resource >= NumResources) {
        throw std::runtime_error("Invalid resource id.");
    }
}

void ResourcesTracker::Add(const Node& node, const Resource& resource,
                           const ResourceQuantity& quantity) {
    ValidateResource(resource);

    auto it = State.find(node);
    if (it != State.end()) {
        nodeState.All[resource] += quantity;
        nodeState.Current[resource] += quantity;
    } else {
        NodeState nodeState(NumResources);
        State.insert(Resources::value_type(node, nodeState));
    }
}

bool ResourcesTracker::Acquire(const Node& node, const Resource& resource,
                               const ResourceQuantity& quantity) {
    ValidateResource(resource);

    auto it = State.find(node);
    if (it == State.end()) {
        return false;
    }

    NodeState& nodeState = it->second;
    assert(nodeState.Current.size() == NumResources);

    if (nodeState.Current[resource] >= quantity) {
        nodeState.Current[resource] -= quantity;
        return true;
    }

    return false;
}

} // Meteor
<commit_msg>Fixed compilation<commit_after>#include "resources_tracker.h"

#include <cassert>
#include <stdexcept>

namespace Meteor {

ResourcesTracker::ResourcesTracker(size_t numResources)
    : NumResources(numResources)
{
}

/// @throw runtime_error
void ResourcesTracker::ValidateResource(const Resource& resource) {
    if (resource >= NumResources) {
        throw std::runtime_error("Invalid resource id.");
    }
}

void ResourcesTracker::Add(const Node& node, const Resource& resource,
                           const ResourceQuantity& quantity) {
    ValidateResource(resource);

    auto it = State.find(node);
    if (it != State.end()) {
        NodeState& nodeState = it->second;
        nodeState.All[resource] += quantity;
        nodeState.Current[resource] += quantity;
    } else {
        NodeState nodeState(NumResources);
        State.insert(Resources::value_type(node, nodeState));
    }
}

bool ResourcesTracker::Acquire(const Node& node, const Resource& resource,
                               const ResourceQuantity& quantity) {
    ValidateResource(resource);

    auto it = State.find(node);
    if (it == State.end()) {
        return false;
    }

    NodeState& nodeState = it->second;
    assert(nodeState.Current.size() == NumResources);

    if (nodeState.Current[resource] >= quantity) {
        nodeState.Current[resource] -= quantity;
        return true;
    }

    return false;
}

} // Meteor
<|endoftext|>


<commit_before>/*
 * Collect statistics of a beng-lb process.
 *
 * author: Max Kellermann <mk@cm4all.com>
 */

#include "bp_stats.hxx"
#include "bp_instance.hxx"
#include "stock/MapStock.hxx"
#include "stock/Stats.hxx"
#include "fb_pool.hxx"
#include "SlicePool.hxx"
#include "tcache.hxx"
#include "http_cache.hxx"
#include "fcache.hxx"
#include "nfs_cache.hxx"
#include "child_manager.hxx"
#include "session_manager.hxx"
#include "AllocatorStats.hxx"
#include "beng-proxy/control.h"
#include "util/ByteOrder.hxx"

void
bp_get_stats(const BpInstance *instance,
             struct beng_control_stats *data)
{
    StockStats tcp_stock_stats = {
        .busy = 0,
        .idle = 0,
    };

    hstock_add_stats(*instance->tcp_stock, tcp_stock_stats);

    const auto tcache_stats = instance->translate_cache != nullptr
        ? translate_cache_get_stats(*instance->translate_cache)
        : AllocatorStats::Zero();
    const auto http_cache_stats = http_cache_get_stats(*instance->http_cache);
    const auto fcache_stats = filter_cache_get_stats(*instance->filter_cache);

    data->incoming_connections = ToBE32(instance->num_connections);
    data->outgoing_connections = ToBE32(tcp_stock_stats.busy
                                               + tcp_stock_stats.idle);
    data->children = ToBE32(child_get_count());
    data->sessions = ToBE32(session_manager_get_count());
    data->http_requests = ToBE64(instance->http_request_counter);
    data->translation_cache_size = ToBE64(tcache_stats.netto_size);
    data->http_cache_size = ToBE64(http_cache_stats.netto_size);
    data->filter_cache_size = ToBE64(fcache_stats.netto_size);

    data->translation_cache_brutto_size = ToBE64(tcache_stats.brutto_size);
    data->http_cache_brutto_size = ToBE64(http_cache_stats.brutto_size);
    data->filter_cache_brutto_size = ToBE64(fcache_stats.brutto_size);

#ifdef HAVE_LIBNFS
    const auto nfs_cache_stats = nfs_cache_get_stats(*instance->nfs_cache);
#else
    const auto nfs_cache_stats = AllocatorStats::Zero();
#endif
    data->nfs_cache_size = ToBE64(nfs_cache_stats.netto_size);
    data->nfs_cache_brutto_size = ToBE64(nfs_cache_stats.brutto_size);

    const auto io_buffers_stats = slice_pool_get_stats(fb_pool_get());
    data->io_buffers_size = ToBE64(io_buffers_stats.netto_size);
    data->io_buffers_brutto_size = ToBE64(io_buffers_stats.brutto_size);

    /* TODO: add stats from all worker processes;  */
}
<commit_msg>bp_stats: use ChildProcessRegistry::GetCount()<commit_after>/*
 * Collect statistics of a beng-lb process.
 *
 * author: Max Kellermann <mk@cm4all.com>
 */

#include "bp_stats.hxx"
#include "bp_instance.hxx"
#include "stock/MapStock.hxx"
#include "stock/Stats.hxx"
#include "fb_pool.hxx"
#include "SlicePool.hxx"
#include "tcache.hxx"
#include "http_cache.hxx"
#include "fcache.hxx"
#include "nfs_cache.hxx"
#include "session_manager.hxx"
#include "AllocatorStats.hxx"
#include "beng-proxy/control.h"
#include "util/ByteOrder.hxx"

void
bp_get_stats(const BpInstance *instance,
             struct beng_control_stats *data)
{
    StockStats tcp_stock_stats = {
        .busy = 0,
        .idle = 0,
    };

    hstock_add_stats(*instance->tcp_stock, tcp_stock_stats);

    const auto tcache_stats = instance->translate_cache != nullptr
        ? translate_cache_get_stats(*instance->translate_cache)
        : AllocatorStats::Zero();
    const auto http_cache_stats = http_cache_get_stats(*instance->http_cache);
    const auto fcache_stats = filter_cache_get_stats(*instance->filter_cache);

    data->incoming_connections = ToBE32(instance->num_connections);
    data->outgoing_connections = ToBE32(tcp_stock_stats.busy
                                               + tcp_stock_stats.idle);
    data->children = ToBE32(instance->child_process_registry.GetCount());
    data->sessions = ToBE32(session_manager_get_count());
    data->http_requests = ToBE64(instance->http_request_counter);
    data->translation_cache_size = ToBE64(tcache_stats.netto_size);
    data->http_cache_size = ToBE64(http_cache_stats.netto_size);
    data->filter_cache_size = ToBE64(fcache_stats.netto_size);

    data->translation_cache_brutto_size = ToBE64(tcache_stats.brutto_size);
    data->http_cache_brutto_size = ToBE64(http_cache_stats.brutto_size);
    data->filter_cache_brutto_size = ToBE64(fcache_stats.brutto_size);

#ifdef HAVE_LIBNFS
    const auto nfs_cache_stats = nfs_cache_get_stats(*instance->nfs_cache);
#else
    const auto nfs_cache_stats = AllocatorStats::Zero();
#endif
    data->nfs_cache_size = ToBE64(nfs_cache_stats.netto_size);
    data->nfs_cache_brutto_size = ToBE64(nfs_cache_stats.brutto_size);

    const auto io_buffers_stats = slice_pool_get_stats(fb_pool_get());
    data->io_buffers_size = ToBE64(io_buffers_stats.netto_size);
    data->io_buffers_brutto_size = ToBE64(io_buffers_stats.brutto_size);

    /* TODO: add stats from all worker processes;  */
}
<|endoftext|>